Function to instantiate DataModule and retrieve dataset from it (DRY)

src/ptbench/data/base_datamodule.py

@@ -92,3 +92,16 @@ class BaseDataModule(pl.LightningDataModule):
             shuffle=False,
             pin_memory=self.pin_memory,
         )
+
+
+def get_dataset_from_module(module, stage, **module_args):
+    """Instantiates a DataModule and retrieves the corresponding dataset.
+
+    Useful when combining multiple datasets.
+    """
+    module_instance = module(**module_args)
+    module_instance.prepare_data()
+    module_instance.setup(stage=stage)
+    dataset = module_instance.dataset
+
+    return dataset

src/ptbench/data/mc_ch/default.py

@@ -8,7 +8,7 @@ from clapper.logging import setup
 from torch.utils.data.dataset import ConcatDataset
 
 from .. import return_subsets
-from ..base_datamodule import BaseDataModule
+from ..base_datamodule import BaseDataModule, get_dataset_from_module
 from ..montgomery.default import datamodule as mc_datamodule
 from ..shenzhen.default import datamodule as ch_datamodule
 
@@ -37,27 +37,16 @@ class DefaultModule(BaseDataModule):
 
     def setup(self, stage: str):
         # Instantiate other datamodules and get their datasets
-        mc_module = mc_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
-
-        mc_module.prepare_data()
-        mc_module.setup(stage=stage)
-        mc = mc_module.dataset
 
-        ch_module = ch_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
+        module_args = {
+            "train_batch_size": self.train_batch_size,
+            "predict_batch_size": self.predict_batch_size,
+            "drop_incomplete_batch": self.drop_incomplete_batch,
+            "multiproc_kwargs": self.multiproc_kwargs,
+        }
 
-        ch_module.prepare_data()
-        ch_module.setup(stage=stage)
-        ch = ch_module.dataset
+        mc = get_dataset_from_module(mc_datamodule, stage, **module_args)
+        ch = get_dataset_from_module(ch_datamodule, stage, **module_args)
 
         # Combine datasets
         self.dataset = {}

src/ptbench/data/mc_ch/fold_0.py

@@ -8,7 +8,7 @@ from clapper.logging import setup
 from torch.utils.data.dataset import ConcatDataset
 
 from .. import return_subsets
-from ..base_datamodule import BaseDataModule
+from ..base_datamodule import BaseDataModule, get_dataset_from_module
 from ..montgomery.fold_0 import datamodule as mc_datamodule
 from ..shenzhen.fold_0 import datamodule as ch_datamodule
 
@@ -37,27 +37,15 @@ class DefaultModule(BaseDataModule):
 
     def setup(self, stage: str):
         # Instantiate other datamodules and get their datasets
-        mc_module = mc_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
-
-        mc_module.prepare_data()
-        mc_module.setup(stage=stage)
-        mc = mc_module.dataset
-
-        ch_module = ch_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
+        module_args = {
+            "train_batch_size": self.train_batch_size,
+            "predict_batch_size": self.predict_batch_size,
+            "drop_incomplete_batch": self.drop_incomplete_batch,
+            "multiproc_kwargs": self.multiproc_kwargs,
+        }
 
-        ch_module.prepare_data()
-        ch_module.setup(stage=stage)
-        ch = ch_module.dataset
+        mc = get_dataset_from_module(mc_datamodule, stage, **module_args)
+        ch = get_dataset_from_module(ch_datamodule, stage, **module_args)
 
         # Combine datasets
         self.dataset = {}

src/ptbench/data/mc_ch/fold_0_rgb.py

@@ -8,7 +8,7 @@ from clapper.logging import setup
 from torch.utils.data.dataset import ConcatDataset
 
 from .. import return_subsets
-from ..base_datamodule import BaseDataModule
+from ..base_datamodule import BaseDataModule, get_dataset_from_module
 from ..montgomery.fold_0_rgb import datamodule as mc_datamodule
 from ..shenzhen.fold_0_rgb import datamodule as ch_datamodule
 
@@ -37,27 +37,15 @@ class DefaultModule(BaseDataModule):
 
     def setup(self, stage: str):
         # Instantiate other datamodules and get their datasets
-        mc_module = mc_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
-
-        mc_module.prepare_data()
-        mc_module.setup(stage=stage)
-        mc = mc_module.dataset
-
-        ch_module = ch_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
+        module_args = {
+            "train_batch_size": self.train_batch_size,
+            "predict_batch_size": self.predict_batch_size,
+            "drop_incomplete_batch": self.drop_incomplete_batch,
+            "multiproc_kwargs": self.multiproc_kwargs,
+        }
 
-        ch_module.prepare_data()
-        ch_module.setup(stage=stage)
-        ch = ch_module.dataset
+        mc = get_dataset_from_module(mc_datamodule, stage, **module_args)
+        ch = get_dataset_from_module(ch_datamodule, stage, **module_args)
 
         # Combine datasets
         self.dataset = {}

src/ptbench/data/mc_ch/fold_1.py

@@ -8,7 +8,7 @@ from clapper.logging import setup
 from torch.utils.data.dataset import ConcatDataset
 
 from .. import return_subsets
-from ..base_datamodule import BaseDataModule
+from ..base_datamodule import BaseDataModule, get_dataset_from_module
 from ..montgomery.fold_1 import datamodule as mc_datamodule
 from ..shenzhen.fold_1 import datamodule as ch_datamodule
 
@@ -37,27 +37,15 @@ class DefaultModule(BaseDataModule):
 
     def setup(self, stage: str):
         # Instantiate other datamodules and get their datasets
-        mc_module = mc_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
-
-        mc_module.prepare_data()
-        mc_module.setup(stage=stage)
-        mc = mc_module.dataset
-
-        ch_module = ch_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
+        module_args = {
+            "train_batch_size": self.train_batch_size,
+            "predict_batch_size": self.predict_batch_size,
+            "drop_incomplete_batch": self.drop_incomplete_batch,
+            "multiproc_kwargs": self.multiproc_kwargs,
+        }
 
-        ch_module.prepare_data()
-        ch_module.setup(stage=stage)
-        ch = ch_module.dataset
+        mc = get_dataset_from_module(mc_datamodule, stage, **module_args)
+        ch = get_dataset_from_module(ch_datamodule, stage, **module_args)
 
         # Combine datasets
         self.dataset = {}

src/ptbench/data/mc_ch/fold_1_rgb.py

@@ -8,7 +8,7 @@ from clapper.logging import setup
 from torch.utils.data.dataset import ConcatDataset
 
 from .. import return_subsets
-from ..base_datamodule import BaseDataModule
+from ..base_datamodule import BaseDataModule, get_dataset_from_module
 from ..montgomery.fold_1_rgb import datamodule as mc_datamodule
 from ..shenzhen.fold_1_rgb import datamodule as ch_datamodule
 
@@ -37,27 +37,15 @@ class DefaultModule(BaseDataModule):
 
     def setup(self, stage: str):
         # Instantiate other datamodules and get their datasets
-        mc_module = mc_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
-
-        mc_module.prepare_data()
-        mc_module.setup(stage=stage)
-        mc = mc_module.dataset
-
-        ch_module = ch_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
+        module_args = {
+            "train_batch_size": self.train_batch_size,
+            "predict_batch_size": self.predict_batch_size,
+            "drop_incomplete_batch": self.drop_incomplete_batch,
+            "multiproc_kwargs": self.multiproc_kwargs,
+        }
 
-        ch_module.prepare_data()
-        ch_module.setup(stage=stage)
-        ch = ch_module.dataset
+        mc = get_dataset_from_module(mc_datamodule, stage, **module_args)
+        ch = get_dataset_from_module(ch_datamodule, stage, **module_args)
 
         # Combine datasets
         self.dataset = {}

src/ptbench/data/mc_ch/fold_2.py

@@ -8,7 +8,7 @@ from clapper.logging import setup
 from torch.utils.data.dataset import ConcatDataset
 
 from .. import return_subsets
-from ..base_datamodule import BaseDataModule
+from ..base_datamodule import BaseDataModule, get_dataset_from_module
 from ..montgomery.fold_2 import datamodule as mc_datamodule
 from ..shenzhen.fold_2 import datamodule as ch_datamodule
 
@@ -37,27 +37,15 @@ class DefaultModule(BaseDataModule):
 
     def setup(self, stage: str):
         # Instantiate other datamodules and get their datasets
-        mc_module = mc_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
-
-        mc_module.prepare_data()
-        mc_module.setup(stage=stage)
-        mc = mc_module.dataset
-
-        ch_module = ch_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
+        module_args = {
+            "train_batch_size": self.train_batch_size,
+            "predict_batch_size": self.predict_batch_size,
+            "drop_incomplete_batch": self.drop_incomplete_batch,
+            "multiproc_kwargs": self.multiproc_kwargs,
+        }
 
-        ch_module.prepare_data()
-        ch_module.setup(stage=stage)
-        ch = ch_module.dataset
+        mc = get_dataset_from_module(mc_datamodule, stage, **module_args)
+        ch = get_dataset_from_module(ch_datamodule, stage, **module_args)
 
         # Combine datasets
         self.dataset = {}

src/ptbench/data/mc_ch/fold_2_rgb.py

@@ -8,7 +8,7 @@ from clapper.logging import setup
 from torch.utils.data.dataset import ConcatDataset
 
 from .. import return_subsets
-from ..base_datamodule import BaseDataModule
+from ..base_datamodule import BaseDataModule, get_dataset_from_module
 from ..montgomery.fold_2_rgb import datamodule as mc_datamodule
 from ..shenzhen.fold_2_rgb import datamodule as ch_datamodule
 
@@ -37,27 +37,15 @@ class DefaultModule(BaseDataModule):
 
     def setup(self, stage: str):
         # Instantiate other datamodules and get their datasets
-        mc_module = mc_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
-
-        mc_module.prepare_data()
-        mc_module.setup(stage=stage)
-        mc = mc_module.dataset
-
-        ch_module = ch_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
+        module_args = {
+            "train_batch_size": self.train_batch_size,
+            "predict_batch_size": self.predict_batch_size,
+            "drop_incomplete_batch": self.drop_incomplete_batch,
+            "multiproc_kwargs": self.multiproc_kwargs,
+        }
 
-        ch_module.prepare_data()
-        ch_module.setup(stage=stage)
-        ch = ch_module.dataset
+        mc = get_dataset_from_module(mc_datamodule, stage, **module_args)
+        ch = get_dataset_from_module(ch_datamodule, stage, **module_args)
 
         # Combine datasets
         self.dataset = {}

src/ptbench/data/mc_ch/fold_3.py

@@ -8,7 +8,7 @@ from clapper.logging import setup
 from torch.utils.data.dataset import ConcatDataset
 
 from .. import return_subsets
-from ..base_datamodule import BaseDataModule
+from ..base_datamodule import BaseDataModule, get_dataset_from_module
 from ..montgomery.fold_3 import datamodule as mc_datamodule
 from ..shenzhen.fold_3 import datamodule as ch_datamodule
 
@@ -37,27 +37,15 @@ class DefaultModule(BaseDataModule):
 
     def setup(self, stage: str):
         # Instantiate other datamodules and get their datasets
-        mc_module = mc_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
-
-        mc_module.prepare_data()
-        mc_module.setup(stage=stage)
-        mc = mc_module.dataset
-
-        ch_module = ch_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
+        module_args = {
+            "train_batch_size": self.train_batch_size,
+            "predict_batch_size": self.predict_batch_size,
+            "drop_incomplete_batch": self.drop_incomplete_batch,
+            "multiproc_kwargs": self.multiproc_kwargs,
+        }
 
-        ch_module.prepare_data()
-        ch_module.setup(stage=stage)
-        ch = ch_module.dataset
+        mc = get_dataset_from_module(mc_datamodule, stage, **module_args)
+        ch = get_dataset_from_module(ch_datamodule, stage, **module_args)
 
         # Combine datasets
         self.dataset = {}

src/ptbench/data/mc_ch/fold_3_rgb.py

@@ -8,7 +8,7 @@ from clapper.logging import setup
 from torch.utils.data.dataset import ConcatDataset
 
 from .. import return_subsets
-from ..base_datamodule import BaseDataModule
+from ..base_datamodule import BaseDataModule, get_dataset_from_module
 from ..montgomery.fold_3_rgb import datamodule as mc_datamodule
 from ..shenzhen.fold_3_rgb import datamodule as ch_datamodule
 
@@ -37,27 +37,15 @@ class DefaultModule(BaseDataModule):
 
     def setup(self, stage: str):
         # Instantiate other datamodules and get their datasets
-        mc_module = mc_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
-
-        mc_module.prepare_data()
-        mc_module.setup(stage=stage)
-        mc = mc_module.dataset
-
-        ch_module = ch_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
+        module_args = {
+            "train_batch_size": self.train_batch_size,
+            "predict_batch_size": self.predict_batch_size,
+            "drop_incomplete_batch": self.drop_incomplete_batch,
+            "multiproc_kwargs": self.multiproc_kwargs,
+        }
 
-        ch_module.prepare_data()
-        ch_module.setup(stage=stage)
-        ch = ch_module.dataset
+        mc = get_dataset_from_module(mc_datamodule, stage, **module_args)
+        ch = get_dataset_from_module(ch_datamodule, stage, **module_args)
 
         # Combine datasets
         self.dataset = {}

src/ptbench/data/mc_ch/fold_4.py

@@ -8,7 +8,7 @@ from clapper.logging import setup
 from torch.utils.data.dataset import ConcatDataset
 
 from .. import return_subsets
-from ..base_datamodule import BaseDataModule
+from ..base_datamodule import BaseDataModule, get_dataset_from_module
 from ..montgomery.fold_4 import datamodule as mc_datamodule
 from ..shenzhen.fold_4 import datamodule as ch_datamodule
 
@@ -37,27 +37,15 @@ class DefaultModule(BaseDataModule):
 
     def setup(self, stage: str):
         # Instantiate other datamodules and get their datasets
-        mc_module = mc_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
-
-        mc_module.prepare_data()
-        mc_module.setup(stage=stage)
-        mc = mc_module.dataset
-
-        ch_module = ch_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
+        module_args = {
+            "train_batch_size": self.train_batch_size,
+            "predict_batch_size": self.predict_batch_size,
+            "drop_incomplete_batch": self.drop_incomplete_batch,
+            "multiproc_kwargs": self.multiproc_kwargs,
+        }
 
-        ch_module.prepare_data()
-        ch_module.setup(stage=stage)
-        ch = ch_module.dataset
+        mc = get_dataset_from_module(mc_datamodule, stage, **module_args)
+        ch = get_dataset_from_module(ch_datamodule, stage, **module_args)
 
         # Combine datasets
         self.dataset = {}

src/ptbench/data/mc_ch/fold_4_rgb.py

@@ -8,7 +8,7 @@ from clapper.logging import setup
 from torch.utils.data.dataset import ConcatDataset
 
 from .. import return_subsets
-from ..base_datamodule import BaseDataModule
+from ..base_datamodule import BaseDataModule, get_dataset_from_module
 from ..montgomery.fold_4_rgb import datamodule as mc_datamodule
 from ..shenzhen.fold_4_rgb import datamodule as ch_datamodule
 
@@ -37,27 +37,15 @@ class DefaultModule(BaseDataModule):
 
     def setup(self, stage: str):
         # Instantiate other datamodules and get their datasets
-        mc_module = mc_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
-
-        mc_module.prepare_data()
-        mc_module.setup(stage=stage)
-        mc = mc_module.dataset
-
-        ch_module = ch_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
+        module_args = {
+            "train_batch_size": self.train_batch_size,
+            "predict_batch_size": self.predict_batch_size,
+            "drop_incomplete_batch": self.drop_incomplete_batch,
+            "multiproc_kwargs": self.multiproc_kwargs,
+        }
 
-        ch_module.prepare_data()
-        ch_module.setup(stage=stage)
-        ch = ch_module.dataset
+        mc = get_dataset_from_module(mc_datamodule, stage, **module_args)
+        ch = get_dataset_from_module(ch_datamodule, stage, **module_args)
 
         # Combine datasets
         self.dataset = {}

src/ptbench/data/mc_ch/fold_5.py

@@ -8,7 +8,7 @@ from clapper.logging import setup
 from torch.utils.data.dataset import ConcatDataset
 
 from .. import return_subsets
-from ..base_datamodule import BaseDataModule
+from ..base_datamodule import BaseDataModule, get_dataset_from_module
 from ..montgomery.fold_5 import datamodule as mc_datamodule
 from ..shenzhen.fold_5 import datamodule as ch_datamodule
 
@@ -37,27 +37,15 @@ class DefaultModule(BaseDataModule):
 
     def setup(self, stage: str):
         # Instantiate other datamodules and get their datasets
-        mc_module = mc_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
-
-        mc_module.prepare_data()
-        mc_module.setup(stage=stage)
-        mc = mc_module.dataset
-
-        ch_module = ch_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
+        module_args = {
+            "train_batch_size": self.train_batch_size,
+            "predict_batch_size": self.predict_batch_size,
+            "drop_incomplete_batch": self.drop_incomplete_batch,
+            "multiproc_kwargs": self.multiproc_kwargs,
+        }
 
-        ch_module.prepare_data()
-        ch_module.setup(stage=stage)
-        ch = ch_module.dataset
+        mc = get_dataset_from_module(mc_datamodule, stage, **module_args)
+        ch = get_dataset_from_module(ch_datamodule, stage, **module_args)
 
         # Combine datasets
         self.dataset = {}

src/ptbench/data/mc_ch/fold_5_rgb.py

@@ -8,7 +8,7 @@ from clapper.logging import setup
 from torch.utils.data.dataset import ConcatDataset
 
 from .. import return_subsets
-from ..base_datamodule import BaseDataModule
+from ..base_datamodule import BaseDataModule, get_dataset_from_module
 from ..montgomery.fold_5_rgb import datamodule as mc_datamodule
 from ..shenzhen.fold_5_rgb import datamodule as ch_datamodule
 
@@ -37,27 +37,15 @@ class DefaultModule(BaseDataModule):
 
     def setup(self, stage: str):
         # Instantiate other datamodules and get their datasets
-        mc_module = mc_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
-
-        mc_module.prepare_data()
-        mc_module.setup(stage=stage)
-        mc = mc_module.dataset
-
-        ch_module = ch_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
+        module_args = {
+            "train_batch_size": self.train_batch_size,
+            "predict_batch_size": self.predict_batch_size,
+            "drop_incomplete_batch": self.drop_incomplete_batch,
+            "multiproc_kwargs": self.multiproc_kwargs,
+        }
 
-        ch_module.prepare_data()
-        ch_module.setup(stage=stage)
-        ch = ch_module.dataset
+        mc = get_dataset_from_module(mc_datamodule, stage, **module_args)
+        ch = get_dataset_from_module(ch_datamodule, stage, **module_args)
 
         # Combine datasets
         self.dataset = {}

src/ptbench/data/mc_ch/fold_6.py

@@ -8,7 +8,7 @@ from clapper.logging import setup
 from torch.utils.data.dataset import ConcatDataset
 
 from .. import return_subsets
-from ..base_datamodule import BaseDataModule
+from ..base_datamodule import BaseDataModule, get_dataset_from_module
 from ..montgomery.fold_6 import datamodule as mc_datamodule
 from ..shenzhen.fold_6 import datamodule as ch_datamodule
 
@@ -37,27 +37,15 @@ class DefaultModule(BaseDataModule):
 
     def setup(self, stage: str):
         # Instantiate other datamodules and get their datasets
-        mc_module = mc_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
-
-        mc_module.prepare_data()
-        mc_module.setup(stage=stage)
-        mc = mc_module.dataset
-
-        ch_module = ch_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
+        module_args = {
+            "train_batch_size": self.train_batch_size,
+            "predict_batch_size": self.predict_batch_size,
+            "drop_incomplete_batch": self.drop_incomplete_batch,
+            "multiproc_kwargs": self.multiproc_kwargs,
+        }
 
-        ch_module.prepare_data()
-        ch_module.setup(stage=stage)
-        ch = ch_module.dataset
+        mc = get_dataset_from_module(mc_datamodule, stage, **module_args)
+        ch = get_dataset_from_module(ch_datamodule, stage, **module_args)
 
         # Combine datasets
         self.dataset = {}

src/ptbench/data/mc_ch/fold_6_rgb.py

@@ -8,7 +8,7 @@ from clapper.logging import setup
 from torch.utils.data.dataset import ConcatDataset
 
 from .. import return_subsets
-from ..base_datamodule import BaseDataModule
+from ..base_datamodule import BaseDataModule, get_dataset_from_module
 from ..montgomery.fold_6_rgb import datamodule as mc_datamodule
 from ..shenzhen.fold_6_rgb import datamodule as ch_datamodule
 
@@ -37,27 +37,15 @@ class DefaultModule(BaseDataModule):
 
     def setup(self, stage: str):
         # Instantiate other datamodules and get their datasets
-        mc_module = mc_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
-
-        mc_module.prepare_data()
-        mc_module.setup(stage=stage)
-        mc = mc_module.dataset
-
-        ch_module = ch_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
+        module_args = {
+            "train_batch_size": self.train_batch_size,
+            "predict_batch_size": self.predict_batch_size,
+            "drop_incomplete_batch": self.drop_incomplete_batch,
+            "multiproc_kwargs": self.multiproc_kwargs,
+        }
 
-        ch_module.prepare_data()
-        ch_module.setup(stage=stage)
-        ch = ch_module.dataset
+        mc = get_dataset_from_module(mc_datamodule, stage, **module_args)
+        ch = get_dataset_from_module(ch_datamodule, stage, **module_args)
 
         # Combine datasets
         self.dataset = {}

src/ptbench/data/mc_ch/fold_7.py

@@ -8,7 +8,7 @@ from clapper.logging import setup
 from torch.utils.data.dataset import ConcatDataset
 
 from .. import return_subsets
-from ..base_datamodule import BaseDataModule
+from ..base_datamodule import BaseDataModule, get_dataset_from_module
 from ..montgomery.fold_7 import datamodule as mc_datamodule
 from ..shenzhen.fold_7 import datamodule as ch_datamodule
 
@@ -37,27 +37,15 @@ class DefaultModule(BaseDataModule):
 
     def setup(self, stage: str):
         # Instantiate other datamodules and get their datasets
-        mc_module = mc_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
-
-        mc_module.prepare_data()
-        mc_module.setup(stage=stage)
-        mc = mc_module.dataset
-
-        ch_module = ch_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
+        module_args = {
+            "train_batch_size": self.train_batch_size,
+            "predict_batch_size": self.predict_batch_size,
+            "drop_incomplete_batch": self.drop_incomplete_batch,
+            "multiproc_kwargs": self.multiproc_kwargs,
+        }
 
-        ch_module.prepare_data()
-        ch_module.setup(stage=stage)
-        ch = ch_module.dataset
+        mc = get_dataset_from_module(mc_datamodule, stage, **module_args)
+        ch = get_dataset_from_module(ch_datamodule, stage, **module_args)
 
         # Combine datasets
         self.dataset = {}

src/ptbench/data/mc_ch/fold_7_rgb.py

@@ -8,7 +8,7 @@ from clapper.logging import setup
 from torch.utils.data.dataset import ConcatDataset
 
 from .. import return_subsets
-from ..base_datamodule import BaseDataModule
+from ..base_datamodule import BaseDataModule, get_dataset_from_module
 from ..montgomery.fold_7_rgb import datamodule as mc_datamodule
 from ..shenzhen.fold_7_rgb import datamodule as ch_datamodule
 
@@ -37,27 +37,15 @@ class DefaultModule(BaseDataModule):
 
     def setup(self, stage: str):
         # Instantiate other datamodules and get their datasets
-        mc_module = mc_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
-
-        mc_module.prepare_data()
-        mc_module.setup(stage=stage)
-        mc = mc_module.dataset
-
-        ch_module = ch_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
+        module_args = {
+            "train_batch_size": self.train_batch_size,
+            "predict_batch_size": self.predict_batch_size,
+            "drop_incomplete_batch": self.drop_incomplete_batch,
+            "multiproc_kwargs": self.multiproc_kwargs,
+        }
 
-        ch_module.prepare_data()
-        ch_module.setup(stage=stage)
-        ch = ch_module.dataset
+        mc = get_dataset_from_module(mc_datamodule, stage, **module_args)
+        ch = get_dataset_from_module(ch_datamodule, stage, **module_args)
 
         # Combine datasets
         self.dataset = {}

src/ptbench/data/mc_ch/fold_8.py

@@ -8,7 +8,7 @@ from clapper.logging import setup
 from torch.utils.data.dataset import ConcatDataset
 
 from .. import return_subsets
-from ..base_datamodule import BaseDataModule
+from ..base_datamodule import BaseDataModule, get_dataset_from_module
 from ..montgomery.fold_8 import datamodule as mc_datamodule
 from ..shenzhen.fold_8 import datamodule as ch_datamodule
 
@@ -37,27 +37,15 @@ class DefaultModule(BaseDataModule):
 
     def setup(self, stage: str):
         # Instantiate other datamodules and get their datasets
-        mc_module = mc_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
-
-        mc_module.prepare_data()
-        mc_module.setup(stage=stage)
-        mc = mc_module.dataset
-
-        ch_module = ch_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
+        module_args = {
+            "train_batch_size": self.train_batch_size,
+            "predict_batch_size": self.predict_batch_size,
+            "drop_incomplete_batch": self.drop_incomplete_batch,
+            "multiproc_kwargs": self.multiproc_kwargs,
+        }
 
-        ch_module.prepare_data()
-        ch_module.setup(stage=stage)
-        ch = ch_module.dataset
+        mc = get_dataset_from_module(mc_datamodule, stage, **module_args)
+        ch = get_dataset_from_module(ch_datamodule, stage, **module_args)
 
         # Combine datasets
         self.dataset = {}

src/ptbench/data/mc_ch/fold_8_rgb.py

@@ -8,7 +8,7 @@ from clapper.logging import setup
 from torch.utils.data.dataset import ConcatDataset
 
 from .. import return_subsets
-from ..base_datamodule import BaseDataModule
+from ..base_datamodule import BaseDataModule, get_dataset_from_module
 from ..montgomery.fold_8_rgb import datamodule as mc_datamodule
 from ..shenzhen.fold_8_rgb import datamodule as ch_datamodule
 
@@ -37,27 +37,15 @@ class DefaultModule(BaseDataModule):
 
     def setup(self, stage: str):
         # Instantiate other datamodules and get their datasets
-        mc_module = mc_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
-
-        mc_module.prepare_data()
-        mc_module.setup(stage=stage)
-        mc = mc_module.dataset
-
-        ch_module = ch_datamodule(
-            train_batch_size=self.train_batch_size,
-            predict_batch_size=self.predict_batch_size,
-            drop_incomplete_batch=self.drop_incomplete_batch,
-            multiproc_kwargs=self.multiproc_kwargs,
-        )
+        module_args = {
+            "train_batch_size": self.train_batch_size,
+            "predict_batch_size": self.predict_batch_size,
+            "drop_incomplete_batch": self.drop_incomplete_batch,
+            "multiproc_kwargs": self.multiproc_kwargs,
+        }
 
-        ch_module.prepare_data()
-        ch_module.setup(stage=stage)
-        ch = ch_module.dataset
+        mc = get_dataset_from_module(mc_datamodule, stage, **module_args)
+        ch = get_dataset_from_module(ch_datamodule, stage, **module_args)
 
         # Combine datasets
         self.dataset = {}