forwarding key:value arguments to the data_loader

Author: Jeremy Fix

train.py

@@ -38,6 +38,15 @@ def train(cfg, writer, logger):
     data_aug = get_composed_augmentations(augmentations)
 
     # Setup Dataloader
+    dataloader_args = cfg["data"].copy()
+    dataloader_args.pop('dataset')
+    dataloader_args.pop('train_split')
+    dataloader_args.pop('val_split')
+    dataloader_args.pop('img_rows')
+    dataloader_args.pop('img_cols')
+    dataloader_args.pop('path')
+
+
     data_loader = get_loader(cfg["data"]["dataset"])
     data_path = cfg["data"]["path"]
 
@@ -47,21 +56,23 @@ def train(cfg, writer, logger):
         split=cfg["data"]["train_split"],
         img_size=(cfg["data"]["img_rows"], cfg["data"]["img_cols"]),
         augmentations=data_aug,
+        **dataloader_args
     )
 
     v_loader = data_loader(
         data_path,
         is_transform=True,
         split=cfg["data"]["val_split"],
         img_size=(cfg["data"]["img_rows"], cfg["data"]["img_cols"]),
+        **dataloader_args
     )
 
     n_classes = t_loader.n_classes
     trainloader = data.DataLoader(
         t_loader,
         batch_size=cfg["training"]["batch_size"],
         num_workers=cfg["training"]["n_workers"],
-        shuffle=True,
+        shuffle=True
     )
 
     valloader = data.DataLoader(