Adding --num_workers input parameter to the EEG_GCNN example.

dmlc · Oct 18, 2023 · 9bda3db · 9bda3db
1 parent 9d5b897
commit 9bda3db
Showing 1 changed file with 10 additions and 4 deletions.
diff --git a/examples/pytorch/eeg-gcnn/main.py b/examples/pytorch/eeg-gcnn/main.py
@@ -37,6 +37,12 @@ def _load_memory_mapped_array(file_name):
     parser.add_argument(
         "--num_nodes", type=int, default=8, help="Number of nodes in the graph"
     )
+    parser.add_argument(
+        "--num_workers",
+        type=int,
+        default=4,
+        help="Number of epochs used to train",
+    )
     parser.add_argument(
         "--gpu_idx",
         type=int,
@@ -149,7 +155,7 @@ def _load_memory_mapped_array(file_name):
     # Dataloader========================================================================================================
 
     # use WeightedRandomSampler to balance the training dataset
-    NUM_WORKERS = 4
+    _NUM_WORKERS = args.num_workers
 
     labels_unique, counts = np.unique(y, return_counts=True)
 
@@ -172,7 +178,7 @@ def _load_memory_mapped_array(file_name):
         dataset=train_dataset,
         batch_size=_BATCH_SIZE,
         sampler=weighted_sampler,
-        num_workers=NUM_WORKERS,
+        num_workers=_NUM_WORKERS,
         pin_memory=True,
     )
 
@@ -181,7 +187,7 @@ def _load_memory_mapped_array(file_name):
         dataset=train_dataset,
         batch_size=_BATCH_SIZE,
         shuffle=False,
-        num_workers=NUM_WORKERS,
+        num_workers=_NUM_WORKERS,
         pin_memory=True,
     )
 
@@ -194,7 +200,7 @@ def _load_memory_mapped_array(file_name):
         dataset=test_dataset,
         batch_size=_BATCH_SIZE,
         shuffle=False,
-        num_workers=NUM_WORKERS,
+        num_workers=_NUM_WORKERS,
         pin_memory=True,
     )