FEAT: load clearml dataset to prepare for training

trainingModel/run_training.py

@@ -0,0 +1,96 @@
+import os
+import numpy as np
+
+from clearml import Task, Dataset
+from datasets import load_dataset
+from dataPrep.helpers.transforms_loaders import make_dataset_loaders
+
+import torch
+from models.modelOne import modelOne
+from trainingModel.Training import train_model
+
+
+# Load data prep task from ClearML
+DATA_PREP_TASK_ID = "f6888baedc7142fcad9e0cc6837c5cb5"
+DATA_PREP = Task.get_task(task_id=DATA_PREP_TASK_ID)
+
+data_params = DATA_PREP.get_parameters()
+dataset_link = data_params['General/dataset/link']
+
+# Load the whole dataset
+try:
+    ds = load_dataset(dataset_link)
+except Exception as e:
+    raise RuntimeError(f"Error loading the dataset: {e}")
+
+full_dataset = ds['train']
+
+
+# Load the subset indices from ClearML
+SUBSET_ID = "f6888baedc7142fcad9e0cc6837c5cb5"
+subset_clearml = Dataset.get(dataset_id=SUBSET_ID)
+
+local_folder = subset_clearml.get_local_copy()
+subset_indices_path = os.path.join(local_folder, "subset_indices.npy")
+subset_indices = np.load(subset_indices_path)
+
+print("Loaded subset indices:", subset_indices.shape)
+
+
+# Apply subset indices to full dataset - this gives you the same subset as data prep
+subset_dataset = full_dataset.select(subset_indices)
+
+
+# Extract parameters from data prep task - these will create the DataLoaders
+seed = int(data_params['General/seed'])
+batch_size = int(data_params['General/dataloaders/batch_size'])
+test_size = float(data_params['General/dataloaders/test_size'])
+
+aug_config = {
+    'rotation': float(data_params['General/augmentation/rotation']),
+    'brightness': float(data_params['General/augmentation/brightness']),
+    'saturation': float(data_params['General/augmentation/saturation']),
+    'blur': float(data_params['General/augmentation/blur'])
+}
+
+# Create DataLoaders using the parameters from data prep
+subset_loaders = make_dataset_loaders(
+    subset_dataset, seed, batch_size, test_size, aug_config
+)
+
+print("\n--- Handoff Test Successful ---")
+print(f"Prototype Train loader batches: {len(subset_loaders['train'])}")
+print(f"Prototype Validation loader batches: {len(subset_loaders['val'])}")
+print(f"Prototype Test loader batches: {len(subset_loaders['test'])}")
+
+
+full_loaders = make_dataset_loaders(
+    full_dataset, seed, batch_size, test_size, aug_config
+)
+
+print("\n--- Handoff Test Successful ---")
+print(f"Train loader batches: {len(full_loaders['train'])}")
+print(f"Validation loader batches: {len(full_loaders['val'])}")
+print(f"Test loader batches: {len(full_loaders['test'])}")
+
+
+# -------- Build the ML model --------
+model = modelOne(noOfClasses=39)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+# ------- Train the model (on subset for now) -------
+''' 
+When calling this function, the model should be trained on the given dataset
+
+
+train_model(
+    model=model,
+    train_loader=subset_loaders['train'],
+    val_loader=subset_loaders['val'],
+    device=device,
+    n_epochs=10,
+    lr=1e-3,
+    save_path="best_model.pt",
+)
+'''