configure dataloader workers

dataPrep/helpers/clearml_data.py

@@ -11,7 +11,7 @@ Takes latest Data Prep ClearML task from project and reconstruct:
 - data loaders for both full and subset datasets
 - Aug settings used
 '''
-def extract_latest_data_task(project_name: str = "Small Group Project"):
+def extract_latest_data_task(project_name: str = "Small Group Project", num_workers: int = 8):
 
   # --------- Get latest Data Preparation task from ClearML ---------
 
@@ -76,7 +76,7 @@ def extract_latest_data_task(project_name: str = "Small Group Project"):
   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)
+  subset_loaders, full_loaders, aug_config = get_data_loaders(data_params, subset_dataset, full_dataset, num_workers=num_workers)
   batch_size = int(data_params['General/dataloaders/batch_size'])
   seed = int(data_params['General/seed'])
 
@@ -99,7 +99,7 @@ def extract_latest_data_task(project_name: str = "Small Group Project"):
 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):
+def get_data_loaders(data_params, subset_dataset, full_dataset, num_workers):
 
   # Extract data parameters- these will be used in the DataLoaders
   seed = int(data_params['General/seed'])
@@ -115,7 +115,7 @@ def get_data_loaders(data_params, subset_dataset, full_dataset):
 
   # Create DataLoaders using the parameters from data prep
   subset_loaders = make_dataset_loaders(
-      subset_dataset, seed, batch_size, test_size, aug_config
+      subset_dataset, seed, batch_size, test_size, aug_config, workers=num_workers
   )
 
   print("\n--- Handoff Test Successful ---")
@@ -125,7 +125,7 @@ def get_data_loaders(data_params, subset_dataset, full_dataset):
 
 
   full_loaders = make_dataset_loaders(
-      full_dataset, seed, batch_size, test_size, aug_config
+      full_dataset, seed, batch_size, test_size, aug_config, workers=num_workers
   )
 
   print("\n--- Handoff Test Successful ---")

dataPrep/helpers/transforms_loaders.py

@@ -47,24 +47,25 @@ def make_augment_pipeline(aug_config):
     return augmentation
 
 
+def apply_augmentation(batch, augmentation):
+    batch['image'] = [augmentation(x) for x in batch['image']]
+    return batch
+
+def apply_normalisation(batch, normalisation):
+    batch['image'] = [normalisation(x) for x in batch['image']]
+    return batch
+    
+
 """
 Creates and returns DataLoaders (train, val, test) for a given dataset.
 Performs a 70/15/15 split
 """
-def make_dataset_loaders(dataset, seed, batch_size, test_size, aug_config):
+def make_dataset_loaders(dataset, seed, batch_size, test_size, aug_config, workers=8):
 
     # Define transformation pipelines for the dataset
     normalisation = make_norm_pipeline()
     augmentation = make_augment_pipeline(aug_config)
 
-    def apply_augmentation(batch):
-        batch['image'] = [augmentation(x) for x in batch['image']]
-        return batch
-
-    def apply_normalisation(batch):
-        batch['image'] = [normalisation(x) for x in batch['image']]
-        return batch
-    
     # 70/30 split creates train set
     split_1 = dataset.train_test_split(test_size=test_size, seed=seed)
     train_split = split_1['train']
@@ -76,14 +77,34 @@ def make_dataset_loaders(dataset, seed, batch_size, test_size, aug_config):
     val_split, test_split = split_2['train'], split_2['test']
 
     # Put each split through pipelines
-    train_split.set_transform(apply_augmentation)
-    val_split.set_transform(apply_normalisation)
-    test_split.set_transform(apply_normalisation)
+    train_split.set_transform(lambda batch: apply_augmentation(batch, augmentation))
+    val_split.set_transform(lambda batch: apply_normalisation(batch, normalisation))
+    test_split.set_transform(lambda batch: apply_normalisation(batch, normalisation))
 
     # Create dataloader for each
-    train_loader = DataLoader(train_split, batch_size=batch_size, shuffle=True)
-    val_loader = DataLoader(val_split, batch_size=batch_size, shuffle=False)
-    test_loader = DataLoader(test_split, batch_size=batch_size, shuffle=False)
+    train_loader = DataLoader(
+        train_split,
+        batch_size=batch_size,
+        shuffle=True,
+        pin_memory=True,
+        num_workers=workers
+    )
+    val_loader = DataLoader(
+        val_split,
+        batch_size=batch_size,
+        shuffle=False,
+        pin_memory=True,
+        num_workers=workers
+    )
+    test_loader = DataLoader(
+        test_split,
+        batch_size=batch_size,
+        shuffle=False,
+        pin_memory=True,
+        num_workers=workers
+    )
+
+    print(f"\nWorkers used in DataLoaders: {workers}\n")
 
     dataset_loaders = {
         "train": train_loader,

trainingModel/run_training.py

@@ -1,4 +1,4 @@
-
+import os
 from clearml import Task
 from dataPrep.helpers.clearml_data import extract_latest_data_task
 
@@ -8,8 +8,9 @@ from trainingModel.helpers.Training import train_model
 
 
 # -------------- Load Data --------------
+NUM_WORKERS = 0
 project_name = "Small Group Project"
-subset_loaders, full_loaders, data_prep_metadata = extract_latest_data_task(project_name=project_name)
+subset_loaders, full_loaders, data_prep_metadata = extract_latest_data_task(project_name=project_name, num_workers=NUM_WORKERS)
 
 
 # -------- ClearML Training Task Setup --------
@@ -30,6 +31,7 @@ training_config = {
     "learning_rate": 1e-3,
     "optimizer": "adam",
     "save_path": "best_model.pt",
+    "num_workers": NUM_WORKERS
 }
 training_task.connect(training_config)