nnUNetTrainerBenchmark_5epochs now saves hostname as well

MIC-DKFZ · May 23, 2024 · e0b33a6 · e0b33a6
1 parent 7372db8
commit e0b33a6
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/nnunetv2/training/nnUNetTrainer/variants/benchmarking/nnUNetTrainerBenchmark_5epochs.py b/nnunetv2/training/nnUNetTrainer/variants/benchmarking/nnUNetTrainerBenchmark_5epochs.py
@@ -1,3 +1,5 @@
+import subprocess
+
 import torch
 from batchgenerators.utilities.file_and_folder_operations import save_json, join, isfile, load_json
 
@@ -53,13 +55,15 @@ def on_train_end(self):
             else:
                 old_results = {}
             # generate some unique key
-            my_key = f"{cudnn_version}__{torch_version.replace(' ', '')}__{gpu_name.replace(' ', '')}__gpus_{num_gpus}"
+            hostname = subprocess.getoutput('hostname')
+            my_key = f"{hostname}__{cudnn_version}__{torch_version.replace(' ', '')}__{gpu_name.replace(' ', '')}__num_gpus_{num_gpus}"
             old_results[my_key] = {
                 'torch_version': torch_version,
                 'cudnn_version': cudnn_version,
                 'gpu_name': gpu_name,
                 'fastest_epoch': fastest_epoch,
                 'num_gpus': num_gpus,
+                'hostname': hostname
             }
             save_json(old_results,
                       join(self.output_folder, 'benchmark_result.json'))