adapt module to dataset refactor

tests/datamodules/test_usavars.py

@@ -10,58 +10,28 @@
 
 
 class TestUSAVarsDataModule:
-    @pytest.fixture(
-        scope="class",
-        params=zip(
-            [["elevation", "population"], ["treecover"]],
-            [True, False],
-            [(0.5, 0.0), (0.0, 0.5)],
-        ),
-    )
+    @pytest.fixture()
     def datamodule(self, request: SubRequest) -> USAVarsDataModule:
-        labels, fixed_shuffle, split = request.param
-        val_split_pct, test_split_pct = split
         root = os.path.join("tests", "data", "usavars")
         batch_size = 1
         num_workers = 0
 
-        dm = USAVarsDataModule(
-            root,
-            labels,
-            None,
-            fixed_shuffle,
-            batch_size,
-            num_workers,
-            val_split_pct,
-            test_split_pct,
-        )
+        dm = USAVarsDataModule(root, batch_size=batch_size, num_workers=num_workers)
         dm.prepare_data()
         dm.setup()
         return dm
 
     def test_train_dataloader(self, datamodule: USAVarsDataModule) -> None:
-        assert len(datamodule.train_dataloader()) == 1
+        assert len(datamodule.train_dataloader()) == 3
         sample = next(iter(datamodule.train_dataloader()))
-        assert sample["labels"].shape[1] == len(datamodule.labels)
-        if datamodule.fixed_shuffle:
-            assert sample["labels"][0, 0] == 1.0
+        assert sample["image"].shape[0] == datamodule.batch_size
 
     def test_val_dataloader(self, datamodule: USAVarsDataModule) -> None:
-        if datamodule.val_split_pct == 0.5:
-            assert len(datamodule.val_dataloader()) == 1
-            sample = next(iter(datamodule.val_dataloader()))
-            assert sample["labels"].shape[1] == len(datamodule.labels)
-            if datamodule.fixed_shuffle:
-                assert sample["labels"][0, 0] == 0.0
-        else:
-            assert len(datamodule.val_dataloader()) == 0
+        assert len(datamodule.val_dataloader()) == 2
+        sample = next(iter(datamodule.val_dataloader()))
+        assert sample["image"].shape[0] == datamodule.batch_size
 
     def test_test_dataloader(self, datamodule: USAVarsDataModule) -> None:
-        if datamodule.test_split_pct == 0.5:
-            assert len(datamodule.test_dataloader()) == 1
-            sample = next(iter(datamodule.test_dataloader()))
-            assert sample["labels"].shape[1] == len(datamodule.labels)
-            if datamodule.fixed_shuffle:
-                assert sample["labels"][0, 0] == 0.0
-        else:
-            assert len(datamodule.test_dataloader()) == 0
+        assert len(datamodule.test_dataloader()) == 1
+        sample = next(iter(datamodule.test_dataloader()))
+        assert sample["image"].shape[0] == datamodule.batch_size

torchgeo/datamodules/usavars.py

@@ -10,7 +10,6 @@
 from torch.utils.data import DataLoader
 
 from ..datasets import USAVars
-from .utils import dataset_split
 
 
 class USAVarsDataModule(pl.LightningModule):
@@ -51,16 +50,21 @@ def prepare_data(self) -> None:
 
         This method is only called once per run.
         """
-        USAVars(self.root_dir, self.labels, checksum=False)
+        USAVars(self.root_dir, labels=self.labels, checksum=False)
 
     def setup(self, stage: Optional[str] = None) -> None:
         """Initialize the main Dataset objects.
 
         This method is called once per GPU per run.
         """
-        dataset = USAVars(self.root_dir, self.labels, transforms=self.transforms)
-        self.train_dataset, self.val_dataset, self.test_dataset = dataset_split(
-            dataset, val_pct=self.val_split_pct, test_pct=self.test_split_pct
+        self.train_dataset = USAVars(
+            self.root_dir, "train", self.labels, transforms=self.transforms
+        )
+        self.val_dataset = USAVars(
+            self.root_dir, "val", self.labels, transforms=self.transforms
+        )
+        self.test_dataset = USAVars(
+            self.root_dir, "test", self.labels, transforms=self.transforms
         )
 
     def train_dataloader(self) -> DataLoader[Any]:
@@ -69,7 +73,7 @@ def train_dataloader(self) -> DataLoader[Any]:
             self.train_dataset,
             batch_size=self.batch_size,
             num_workers=self.num_workers,
-            shuffle=True,
+            shuffle=False,
         )
 
     def val_dataloader(self) -> DataLoader[Any]: