1+
2+ # task.connect(args)
3+ #
4+ # # only create the task, we will actually execute it later
5+ # # task.execute_remotely() # After passing local testing, you should uncomment this command to initial task to ClearML
6+ #
7+ # print('Retrieving Iris dataset')
8+ # dataset_task = Task.get_task(task_id=args['dataset_task_id'])
9+ # X_train = dataset_task.artifacts['X_train'].get()
10+ # X_test = dataset_task.artifacts['X_test'].get()
11+ # y_train = dataset_task.artifacts['y_train'].get()
12+ # y_test = dataset_task.artifacts['y_test'].get()
13+ # print('Iris dataset loaded')
14+ #
15+ #
16+ # # Define a simple neural network
17+ # class SimpleNN(nn.Module):
18+ # def __init__(self, input_size, num_classes):
19+ # super(SimpleNN, self).__init__()
20+ # self.fc1 = nn.Linear(input_size, 50)
21+ # self.fc2 = nn.Linear(50, num_classes)
22+ #
23+ # def forward(self, x):
24+ # x = torch.relu(self.fc1(x))
25+ # x = self.fc2(x)
26+ # return x
27+ #
28+ # # Convert data to PyTorch tensors
29+ # X_train_tensor = torch.tensor(X_train, dtype=torch.float32)
30+ # y_train_tensor = torch.tensor(y_train, dtype=torch.long)
31+ # X_test_tensor = torch.tensor(X_test, dtype=torch.float32)
32+ # y_test_tensor = torch.tensor(y_test, dtype=torch.long)
33+ #
34+ # # Create DataLoader
35+ # train_dataset = TensorDataset(X_train_tensor, y_train_tensor)
36+ # train_loader = DataLoader(train_dataset, batch_size=args['batch_size'], shuffle=True)
37+ # # Hyperparameters
38+ # # Initialize the model, loss function, and optimizer
39+ # model = SimpleNN(input_size=X_train.shape[1], num_classes=len(set(y_train)))
40+ # criterion = nn.CrossEntropyLoss()
41+ # optimizer = optim.Adam(
42+ # model.parameters(),
43+ # lr=args['learning_rate'],
44+ # weight_decay=args['weight_decay']
45+ # )
46+ #
47+ # for epoch in tqdm(range(args['num_epochs']), desc='Training Epochs'):
48+ # epoch_loss = 0.0
49+ #
50+ # for inputs, labels in train_loader:
51+ # optimizer.zero_grad()
52+ # outputs = model(inputs)
53+ # loss = criterion(outputs, labels)
54+ # loss.backward()
55+ # optimizer.step()
56+ #
57+ # # 累积 loss
58+ # epoch_loss += loss.item()
59+ #
60+ # avg_loss = epoch_loss / len(train_loader)
61+ # logger.report_scalar(title='train', series='epoch_loss', value=avg_loss, iteration=epoch)
62+ #
63+ # # Save model
64+ # model_path = 'assets/model.pkl'
65+ # torch.save(model.state_dict(), model_path)
66+ # task.upload_artifact(name='model', artifact_object=model_path)
67+ # print('Model saved and uploaded as artifact')
68+ #
69+ # # Load model for evaluation
70+ # model.load_state_dict(torch.load(model_path))
71+ # model.eval()
72+ # with torch.no_grad():
73+ # outputs = model(X_test_tensor)
74+ # _, predicted = torch.max(outputs, 1)
75+ # accuracy = (predicted == y_test_tensor).float().mean().item()
76+ # logger.report_scalar("validation_accuracy", "score", value=accuracy, iteration=0)
77+ #
78+ # print(f'Model trained & stored with accuracy: {accuracy:.4f}')
79+ #
80+ #
81+ # # Plotting confusion matrix
82+ # species_mapping = {0: 'Setosa', 1: 'Versicolor', 2: 'Virginica'}
83+ # y_test_names = [species_mapping[label.item()] for label in y_test]
84+ # predicted_names = [species_mapping[label.item()] for label in predicted]
85+ #
86+ # cm = confusion_matrix(y_test_names, predicted_names, labels=list(species_mapping.values()))
87+ # disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=list(species_mapping.values()))
88+ # disp.plot(cmap=plt.cm.Blues)
89+ #
90+ # plt.title('Confusion Matrix')
91+ # plt.savefig('figs/confusion_matrix.png')
92+ #
93+ # print('Confusion matrix plotted and saved as confusion_matrix.png')
94+
195import matplotlib .pyplot as plt
96+ import numpy as np
297from clearml import Task , Logger
398import torch
499import torch .nn as nn
5100import torch .optim as optim
6101from torch .utils .data import DataLoader , TensorDataset
7- from sklearn .metrics import confusion_matrix , ConfusionMatrixDisplay
8- from tqdm import tqdm
102+
9103
10104# Connecting ClearML with the current process,
105+ # from here on everything is logged automatically
11106task = Task .init (project_name = "AI_Studio_Demo" , task_name = "Pipeline step 3 train model" )
12107logger = Logger .current_logger ()
13108
109+ # Arguments
14110args = {
15- 'dataset_task_id' : '86f09f66d88c4ec7819f05b086deea15' , # update id if it needs running locally
16- 'num_epochs' : 20 ,
17- 'batch_size' : 16 ,
18- 'dataset_task_id' : '' ,
19-
20- # ✅ HPO
21- 'learning_rate' : 1e-3 ,
22- 'weight_decay' : 1e-5 ,
111+ 'dataset_task_id' : '' , # replace the value only when you need debug locally
23112}
24-
25113task .connect (args )
26114
27115# only create the task, we will actually execute it later
28- # task.execute_remotely() # After passing local testing, you should uncomment this command to initial task to ClearML
29-
116+ task .execute_remotely () # After passing local testing, you should uncomment this command to initial task to ClearML
30117
31118print ('Retrieving Iris dataset' )
32119dataset_task = Task .get_task (task_id = args ['dataset_task_id' ])
36123y_test = dataset_task .artifacts ['y_test' ].get ()
37124print ('Iris dataset loaded' )
38125
39-
40126# Define a simple neural network
41127class SimpleNN (nn .Module ):
42128 def __init__ (self , input_size , num_classes ):
@@ -57,61 +143,51 @@ def forward(self, x):
57143
58144# Create DataLoader
59145train_dataset = TensorDataset (X_train_tensor , y_train_tensor )
60- train_loader = DataLoader (train_dataset , batch_size = args [ 'batch_size' ] , shuffle = True )
61- # Hyperparameters
146+ train_loader = DataLoader (train_dataset , batch_size = 32 , shuffle = True )
147+
62148# Initialize the model, loss function, and optimizer
63149model = SimpleNN (input_size = X_train .shape [1 ], num_classes = len (set (y_train )))
64150criterion = nn .CrossEntropyLoss ()
65- optimizer = optim .Adam (
66- model .parameters (),
67- lr = args ['learning_rate' ],
68- weight_decay = args ['weight_decay' ]
69- )
70-
71- for epoch in tqdm (range (args ['num_epochs' ]), desc = 'Training Epochs' ):
72- epoch_loss = 0.0
151+ optimizer = optim .Adam (model .parameters (), lr = 0.001 )
73152
153+ # Train the model
154+ num_epochs = 20
155+ for epoch in range (num_epochs ):
74156 for inputs , labels in train_loader :
75157 optimizer .zero_grad ()
76158 outputs = model (inputs )
77159 loss = criterion (outputs , labels )
78160 loss .backward ()
79161 optimizer .step ()
162+ logger .report_scalar (title = 'train' , series = 'loss' , value = loss .item (), iteration = epoch )
80163
81- # 累积 loss
82- epoch_loss += loss .item ()
83-
84- avg_loss = epoch_loss / len (train_loader )
85- logger .report_scalar (title = 'train' , series = 'epoch_loss' , value = avg_loss , iteration = epoch )
86-
87- # Save model
88- model_path = 'assets/model.pkl'
89- torch .save (model .state_dict (), model_path )
90- task .upload_artifact (name = 'model' , artifact_object = model_path )
91- print ('Model saved and uploaded as artifact' )
92-
93- # Load model for evaluation
94- model .load_state_dict (torch .load (model_path ))
164+ # Evaluate the model
95165model .eval ()
96166with torch .no_grad ():
97167 outputs = model (X_test_tensor )
98168 _ , predicted = torch .max (outputs , 1 )
99169 accuracy = (predicted == y_test_tensor ).float ().mean ().item ()
100- logger .report_scalar ("validation_accuracy" , "score" , value = accuracy , iteration = 0 )
101170
102171print (f'Model trained & stored with accuracy: { accuracy :.4f} ' )
103172
173+ # Plotting (same as before)
174+ x_min , x_max = X_test [:, 0 ].min () - .5 , X_test [:, 0 ].max () + .5
175+ y_min , y_max = X_test [:, 1 ].min () - .5 , X_test [:, 1 ].max () + .5
176+ h = .02 # step size in the mesh
177+ xx , yy = np .meshgrid (np .arange (x_min , x_max , h ), np .arange (y_min , y_max , h ))
178+ plt .figure (1 , figsize = (4 , 3 ))
179+
104180
105- # Plotting confusion matrix
106- species_mapping = {0 : 'Setosa' , 1 : 'Versicolor' , 2 : 'Virginica' }
107- y_test_names = [species_mapping [label .item ()] for label in y_test ]
108- predicted_names = [species_mapping [label .item ()] for label in predicted ]
181+ plt .scatter (X_test [:, 0 ], X_test [:, 1 ], c = y_test , edgecolors = 'k' , cmap = plt .cm .Paired )
182+ plt .xlabel ('Sepal length' )
183+ plt .ylabel ('Sepal width' )
109184
110- cm = confusion_matrix (y_test_names , predicted_names , labels = list (species_mapping .values ()))
111- disp = ConfusionMatrixDisplay (confusion_matrix = cm , display_labels = list (species_mapping .values ()))
112- disp .plot (cmap = plt .cm .Blues )
185+ plt .xlim (xx .min (), xx .max ())
186+ plt .ylim (yy .min (), yy .max ())
187+ plt .xticks (())
188+ plt .yticks (())
113189
114- plt .title ('Confusion Matrix ' )
115- plt .savefig ('figs/confusion_matrix .png' )
190+ plt .title ('Iris Types ' )
191+ plt .savefig ('iris_plot .png' )
116192
117- print ('Confusion matrix plotted and saved as confusion_matrix.png ' )
193+ print ('Done🔥 ' )
0 commit comments