Update run_workflow.py

uzh-dqbm-cmi · Jun 10, 2021 · 66b4983 · 66b4983
1 parent 343e2e7
commit 66b4983
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/ddi/run_workflow.py b/ddi/run_workflow.py
@@ -718,11 +718,18 @@ def test_run(datatensor_partitions,
         else:
             print('WARNING: train dir not found: {}'.format(path))
 
-def train_test_partition(datatensor_partition, config_map, tr_val_dir, fold_gpu_map):
+# def train_test_partition(datatensor_partition, config_map, tr_val_dir, fold_gpu_map):
+#     config_epochs = config_map[0]['model_config'].num_epochs
+#     print(config_epochs)
+#     train_val_run(datatensor_partition, config_map, tr_val_dir, fold_gpu_map, num_epochs=config_epochs)
+#     test_run(datatensor_partition, config_map, tr_val_dir, tr_val_dir, fold_gpu_map, num_epochs=1)
+
+def train_test_partition(gpu_num, datatensor_partition, config_map, tr_val_dir, fold_gpu_map, queue):
     config_epochs = config_map[0]['model_config'].num_epochs
     print(config_epochs)
     train_val_run(datatensor_partition, config_map, tr_val_dir, fold_gpu_map, num_epochs=config_epochs)
     test_run(datatensor_partition, config_map, tr_val_dir, tr_val_dir, fold_gpu_map, num_epochs=1)
+    queue.put(gpu_num)
 
 def test_partition(datatensor_partition, config_map, tr_val_dir, fold_gpu_map, suffix_testfname):
     config_epochs = config_map[0]['model_config'].num_epochs