chore: extract load data prep from training

dataPrep/helpers/clearml_data.py

@@ -0,0 +1,110 @@
+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
+
+
+'''
+Takes latest Data Prep ClearML task from project and extracts data loaders and metadata
+'''
+def extract_latest_data_task(project_name: str = "Small Group Project"):
+
+  all_tasks = Task.get_tasks(project_name=project_name)
+  if not all_tasks:
+      raise RuntimeError(f"No tasks found in project '{project_name}'")
+
+  dp_tasks = [t for t in all_tasks if t.name == "Data Preparation"]
+  if not dp_tasks:
+      raise RuntimeError("No 'Data Preparation' tasks found in this project!")
+
+  # Latest Data Prep Task 
+  latest_task = max(dp_tasks, key=lambda t: t.id)
+  DYNAMIC_TASK_ID = latest_task.id
+  DATA_PREP = Task.get_task(task_id=DYNAMIC_TASK_ID)
+
+  # Dataset ID 
+  config_objects = DATA_PREP.get_configuration_objects()
+  raw_meta = config_objects["Dataset Metadata"]
+  dataset_id = raw_meta.split("=")[1].strip().replace('"', "")
+
+  # Load ClearML Dataset
+  subset_clearml = Dataset.get(dataset_id=dataset_id)
+  local_folder = subset_clearml.get_local_copy()
+
+  subset_indices = np.load(os.path.join(local_folder, "subset_indices.npy"))
+
+  # Load Dataset Parameters 
+  data_params = DATA_PREP.get_parameters()
+  dataset_link = data_params['General/dataset/link']
+
+  # Load Full Dataset 
+  try:
+      ds = load_dataset(dataset_link)
+  except Exception as e:
+      raise RuntimeError(f"Error loading the dataset: {e}")
+
+  full_dataset = ds['train']
+
+  # Apply subset indices to full dataset - this gives you the same subset as data prep
+  subset_dataset = full_dataset.select(subset_indices)
+
+  # Get data loaders for both full and subset datasets
+  subset_loaders, full_loaders, aug_config = get_data_loaders(data_params, subset_dataset, full_dataset)
+  batch_size = int(data_params['General/dataloaders/batch_size'])
+  seed = int(data_params['General/seed'])
+
+
+  # Gather data prep task metadata
+  data_prep_metadata = {
+    "data_prep_task_id": DYNAMIC_TASK_ID,
+    "dataset_id": dataset_id,
+    "dataset_link": dataset_link,
+    "augmentation_used": aug_config,
+    "batch_size_used": batch_size,
+    "seed_used": seed,
+  }
+
+  return subset_loaders, full_loaders, data_prep_metadata
+
+
+'''
+Takes a given dataset, subset, data params to create DataLoaders
+Loaders split data into train, val, test
+'''
+def get_data_loaders(data_params, subset_dataset, full_dataset):
+  
+  # Extract data parameters- these will be used in 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'])}")
+
+  return subset_loaders, full_loaders, aug_config

trainingModel/run_training.py

@@ -1,9 +1,6 @@
-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
+from clearml import Task
+from dataPrep.helpers.clearml_data import extract_latest_data_task
 
 import torch
 from models.modelOne import modelOne
@@ -11,79 +8,8 @@ from trainingModel.Training import train_model
 
 
 # -------------- Load Data --------------
-
-all_tasks = Task.get_tasks(project_name="Small Group Project")
-if not all_tasks:
-    raise RuntimeError("No tasks found in project 'Small Group Project'")
-
-dp_tasks = [t for t in all_tasks if t.name == "Data Preparation"]
-if not dp_tasks:
-    raise RuntimeError("No 'Data Preparation' tasks found in this project!")
-
-# Latest Data Prep Task 
-latest_task = max(dp_tasks, key=lambda t: t.id)
-DYNAMIC_TASK_ID = latest_task.id
-DATA_PREP = Task.get_task(task_id=DYNAMIC_TASK_ID)
-
-# Dataset ID 
-config_objects = DATA_PREP.get_configuration_objects()
-raw_meta = config_objects["Dataset Metadata"]
-dataset_id = raw_meta.split("=")[1].strip().replace('"', "")
-
-# Load ClearML Dataset
-subset_clearml = Dataset.get(dataset_id=dataset_id)
-local_folder = subset_clearml.get_local_copy()
-
-subset_indices = np.load(os.path.join(local_folder, "subset_indices.npy"))
-
-# Load Dataset Parameters 
-data_params = DATA_PREP.get_parameters()
-dataset_link = data_params['General/dataset/link']
-
-# Load Full Dataset 
-try:
-    ds = load_dataset(dataset_link)
-except Exception as e:
-    raise RuntimeError(f"Error loading the dataset: {e}")
-
-full_dataset = ds['train']
-
-# 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'])}")
-# -------------- DATA PREP ENDS --------------
+project_name = "Small Group Project"
+subset_loaders, full_loaders, data_prep_metadata = extract_latest_data_task(project_name=project_name)
 
 
 # -------- ClearML Training Task Setup --------
@@ -95,21 +21,13 @@ training_task = Task.init(
 
 # Detail the data prep task used
 training_logger = training_task.get_logger()
-data_prep_metadata = {
-    "data_prep_task_id": DYNAMIC_TASK_ID,
-    "dataset_id": dataset_id,
-    "dataset_link": dataset_link,
-    "augmentation_used": aug_config,
-    "seed_used": seed,
-}
-training_task.connect(data_prep_metadata, name="data_prep_metadata")
+training_task.connect(data_prep_metadata, name="data_prep_metadata_READONLY")
 
 # Training parameters - Modify these to experiment
 training_config = {
     "num_classes": 39,
     "n_epochs": 3,
     "learning_rate": 1e-3,
-    "batch_size": batch_size,
     "optimizer": "adam",
     "save_path": "best_model.pt",
 }