Enable exporting models with subset datasets to modelzoo

constantinpape · Dec 19, 2023 · 7d91a1b · 7d91a1b
1 parent b4401c6
commit 7d91a1b
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/torch_em/util/modelzoo.py b/torch_em/util/modelzoo.py
@@ -107,7 +107,10 @@ def _get_model(trainer, postprocessing):
 
 def _pad(input_data, trainer):
     try:
-        ndim = trainer.train_loader.dataset.ndim
+        if isinstance(trainer.train_loader.dataset, torch.utils.data.dataset.Subset):
+            ndim = trainer.train_loader.dataset.dataset.ndim
+        else:
+            ndim = trainer.train_loader.dataset.ndim
     except AttributeError:
         ndim = trainer.train_loader.dataset.datasets[0].ndim
     target_dims = ndim + 2
@@ -305,7 +308,10 @@ def _write_weights(model, export_folder):
 
 def _get_preprocessing(trainer):
     try:
-        ndim = trainer.train_loader.dataset.ndim
+        if isinstance(trainer.train_loader.dataset, torch.utils.data.dataset.Subset):
+            ndim = trainer.train_loader.dataset.dataset.ndim
+        else:
+            ndim = trainer.train_loader.dataset.ndim
     except AttributeError:
         ndim = trainer.train_loader.dataset.datasets[0].ndim
     normalizer = get_normalizer(trainer)

diff --git a/torch_em/util/util.py b/torch_em/util/util.py
@@ -196,6 +196,10 @@ def get_normalizer(trainer):
         isinstance(dataset, torch.utils.data.dataset.ConcatDataset)
     ):
         dataset = dataset.datasets[0]
+
+    if isinstance(dataset, torch.utils.data.dataset.Subset):
+        dataset = dataset.dataset
+
     preprocessor = dataset.raw_transform
 
     if hasattr(preprocessor, "normalizer"):