Add ClearML for mlp

models/mlp/train.py

@@ -1,5 +1,5 @@
 import json
-import os
+import os, sys
 import random
 
 import numpy as np
@@ -7,6 +7,8 @@ import torch
 import torch.nn as nn
 import torch.optim as optim
 
+from clearml import Task
+
 from models.prepare_dataset import get_dataloaders
 
 CFG = {
@@ -24,6 +26,21 @@ CFG = {
 }
 
 
+# ==== ClearML Initialisation =============================================
+task = Task.init(
+    project_name="Focus Guard",
+    task_name=f"MLP Model Training",
+    tags=["training", "mlp_model"]
+)
+
+prefix = 'checkpoints/'+task.name+'_'+task.id+'/'
+os.makedirs(prefix, exist_ok=True)
+
+task.connect(CFG)   
+
+
+
+# ==== Model =============================================
 def set_seed(seed: int):
     random.seed(seed)
     np.random.seed(seed)
@@ -152,6 +169,16 @@ def main():
         history["val_loss"].append(round(val_loss, 4))
         history["val_acc"].append(round(val_acc, 4))
 
+
+        # Log scalars to ClearML
+        current_lr = optimizer.param_groups[0]['lr']
+        task.logger.report_scalar("Loss", "Train", float(train_loss), iteration=epoch)
+        task.logger.report_scalar("Accuracy", "Train", float(train_acc), iteration=epoch)    
+        task.logger.report_scalar("Loss", "Val", float(val_loss), iteration=epoch)
+        task.logger.report_scalar("Accuracy", "Val", float(val_acc), iteration=epoch)
+        task.logger.report_scalar("Learning Rate", "LR", float(current_lr), iteration=epoch)
+        task.logger.flush()
+
         marker = ""
         if val_acc > best_val_acc:
             best_val_acc = val_acc

requirements.txt

@@ -5,3 +5,4 @@ torch>=2.0.0
 torchvision>=0.15.0
 scikit-learn>=1.2.0
 joblib>=1.2.0
+clearml>=2.0.2