fix broken network architecture plotting

MIC-DKFZ · Mar 11, 2022 · 7999450 · 7999450
1 parent e49475d
commit 7999450
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/nnunetv2/training/nnunet_modules/nnUNetModule.py b/nnunetv2/training/nnunet_modules/nnUNetModule.py
@@ -125,6 +125,8 @@ def __init__(self, dataset_name_or_id: Union[int, str], plans_name: str, configu
             'train_losses': list(),
             'val_losses': list(),
             'lrs': list(),
+            'epoch_start_timestamps': list(),
+            'epoch_end_timestamps': list()
         }
         # shut up, this logging is great
 
@@ -146,7 +148,8 @@ def plot_network_architecture(self):
             import hiddenlayer as hl
             g = hl.build_graph(self.network,
                                torch.rand((1, len(self.dataset_json["modality"]),
-                                           *self.plans['configurations'][self.configuration]['patch_size'])),
+                                           *self.plans['configurations'][self.configuration]['patch_size']),
+                                          device=self.device),
                                transforms=None)
             g.save(join(self.output_folder, "network_architecture.pdf"))
             del g

diff --git a/setup.py b/setup.py
@@ -31,6 +31,7 @@
               'nnUNetv2_extract_fingerprint = nnunetv2.experiment_planning.plan_and_preprocess:extract_fingerprint',
               'nnUNetv2_plan_experiment = nnunetv2.experiment_planning.plan_and_preprocess:plan_experiment',
               'nnUNetv2_preprocess = nnunetv2.experiment_planning.plan_and_preprocess:preprocess',
+              'nnUNetv2_train = nnunetv2.run.train:nnUNet_train_from_args'
           ],
       },
       keywords=['deep learning', 'image segmentation', 'medical image analysis',