DataModules: pass kwargs directly to datasets (microsoft#730)

* Datamodules: pass kwargs directly to datasets * Rename root_dir -> root in config files * Fix datamodule tests * Fix mypy * Fix tutorial * Specify all kwarg keys * Fix bands vs. band_set * root_dir -> root * Document **kwargs
yichiac · Oct 1, 2022 · f79b6ec · f79b6ec
1 parent baccd45
commit f79b6ec
Show file tree

Hide file tree

Showing 78 changed files with 285 additions and 343 deletions.
diff --git a/README.md b/README.md
@@ -123,7 +123,7 @@ All TorchGeo datasets are compatible with PyTorch data loaders, making them easy
 In order to facilitate direct comparisons between results published in the literature and further reduce the boilerplate code needed to run experiments with datasets in TorchGeo, we have created PyTorch Lightning [*datamodules*](https://torchgeo.readthedocs.io/en/stable/api/datamodules.html) with well-defined train-val-test splits and [*trainers*](https://torchgeo.readthedocs.io/en/stable/api/trainers.html) for various tasks like classification, regression, and semantic segmentation. These datamodules show how to incorporate augmentations from the kornia library, include preprocessing transforms (with pre-calculated channel statistics), and let users easily experiment with hyperparameters related to the data itself (as opposed to the modeling process). Training a semantic segmentation model on the [Inria Aerial Image Labeling](https://project.inria.fr/aerialimagelabeling/) dataset is as easy as a few imports and four lines of code.
 
 ```python
-datamodule = InriaAerialImageLabelingDataModule(root_dir="...", batch_size=64, num_workers=6)
+datamodule = InriaAerialImageLabelingDataModule(root="...", batch_size=64, num_workers=6)
 task = SemanticSegmentationTask(segmentation_model="unet", encoder_weights="imagenet", learning_rate=0.1)
 trainer = Trainer(gpus=1, default_root_dir="...")
 

diff --git a/conf/bigearthnet.yaml b/conf/bigearthnet.yaml
@@ -14,7 +14,7 @@ experiment:
     in_channels: 14
     num_classes: 19
   datamodule:
-    root_dir: "data/bigearthnet"
+    root: "data/bigearthnet"
     bands: "all"
     num_classes: ${experiment.module.num_classes}
     batch_size: 128

diff --git a/conf/byol.yaml b/conf/byol.yaml
@@ -14,7 +14,7 @@ experiment:
     learning_rate: 1e-3
     learning_rate_schedule_patience: 6
   datamodule:
-    root_dir: "data/chesapeake/cvpr"
+    root: "data/chesapeake/cvpr"
     train_splits:
       - "de-train"
     val_splits:

diff --git a/conf/chesapeake_cvpr.yaml b/conf/chesapeake_cvpr.yaml
@@ -20,7 +20,7 @@ experiment:
     ignore_index: null
     imagenet_pretraining: True
   datamodule:
-    root_dir: "data/chesapeake/cvpr"
+    root: "data/chesapeake/cvpr"
     train_splits:
       - "de-train"
     val_splits:

diff --git a/conf/cowc_counting.yaml b/conf/cowc_counting.yaml
@@ -10,7 +10,7 @@ experiment:
     learning_rate_schedule_patience: 2
     pretrained: True
   datamodule:
-    root_dir: "data/cowc_counting"
+    root: "data/cowc_counting"
     seed: 0
     batch_size: 64
     num_workers: 4
diff --git a/conf/cyclone.yaml b/conf/cyclone.yaml
@@ -10,7 +10,7 @@ experiment:
     learning_rate_schedule_patience: 2
     pretrained: True
   datamodule:
-    root_dir: "data/cyclone"
+    root: "data/cyclone"
     seed: 0
     batch_size: 32
     num_workers: 4
diff --git a/conf/etci2021.yaml b/conf/etci2021.yaml
@@ -11,6 +11,6 @@ experiment:
     num_classes: 2
     ignore_index: 0
   datamodule:
-    root_dir: "data/etci2021"
+    root: "data/etci2021"
     batch_size: 32
     num_workers: 4
diff --git a/conf/eurosat.yaml b/conf/eurosat.yaml
@@ -9,6 +9,6 @@ experiment:
     in_channels: 13
     num_classes: 10
   datamodule:
-    root_dir: "data/eurosat"
+    root: "data/eurosat"
     batch_size: 128
     num_workers: 4
diff --git a/conf/inria.yaml b/conf/inria.yaml
@@ -23,7 +23,7 @@ experiment:
     num_classes: 2
     ignore_index: null
   datamodule:
-      root_dir: "data/inria"
+      root: "data/inria"
       batch_size: 2
       num_workers: 32
       patch_size: 512

diff --git a/conf/landcoverai.yaml b/conf/landcoverai.yaml
@@ -17,6 +17,6 @@ experiment:
     num_filters: 256
     ignore_index: null
   datamodule:
-    root_dir: "data/landcoverai"
+    root: "data/landcoverai"
     batch_size: 32
     num_workers: 4
diff --git a/conf/naipchesapeake.yaml b/conf/naipchesapeake.yaml
@@ -18,8 +18,8 @@ experiment:
     num_filters: 64
     ignore_index: null
   datamodule:
-    naip_root_dir: "data/naip"
-    chesapeake_root_dir: "data/chesapeake/BAYWIDE"
+    naip_root: "data/naip"
+    chesapeake_root: "data/chesapeake/BAYWIDE"
     batch_size: 32
     num_workers: 4
     patch_size: 32
diff --git a/conf/oscd.yaml b/conf/oscd.yaml
@@ -18,8 +18,8 @@ experiment:
     num_filters: 256
     ignore_index: 0
   datamodule:
-    root_dir: "data/oscd"
-    batch_size: 32
+    root: "data/oscd"
+    train_batch_size: 32
     num_workers: 4
     val_split_pct: 0.1
     bands: "all"

diff --git a/conf/resisc45.yaml b/conf/resisc45.yaml
@@ -14,6 +14,6 @@ experiment:
     in_channels: 3
     num_classes: 45
   datamodule:
-    root_dir: "data/resisc45"
+    root: "data/resisc45"
     batch_size: 128
     num_workers: 4
diff --git a/conf/sen12ms.yaml b/conf/sen12ms.yaml
@@ -17,7 +17,7 @@ experiment:
     num_classes: 11
     ignore_index: null
   datamodule:
-    root_dir: "data/sen12ms"
+    root: "data/sen12ms"
     band_set: "all"
     batch_size: 32
     num_workers: 4

diff --git a/conf/so2sat.yaml b/conf/so2sat.yaml
@@ -14,8 +14,8 @@ experiment:
     in_channels: 3
     num_classes: 17
   datamodule:
-    root_dir: "data/so2sat"
+    root: "data/so2sat"
     batch_size: 128
     num_workers: 4
-    bands: "rgb"
+    band_set: "rgb"
     unsupervised_mode: False
diff --git a/conf/ucmerced.yaml b/conf/ucmerced.yaml
@@ -9,6 +9,6 @@ experiment:
     in_channels: 3
     num_classes: 21
   datamodule:
-    root_dir: "data/ucmerced"
+    root: "data/ucmerced"
     batch_size: 128
     num_workers: 4
diff --git a/docs/tutorials/trainers.ipynb b/docs/tutorials/trainers.ipynb
@@ -132,7 +132,7 @@
     "data_dir = os.path.join(tempfile.gettempdir(), \"cyclone_data\")\n",
     "\n",
     "datamodule = CycloneDataModule(\n",
-    "    root_dir=data_dir, seed=1337, batch_size=64, num_workers=6, api_key=MLHUB_API_KEY\n",
+    "    root=data_dir, seed=1337, batch_size=64, num_workers=6, api_key=MLHUB_API_KEY\n",
     ")"
    ]
   },

diff --git a/evaluate.py b/evaluate.py
@@ -44,7 +44,7 @@ def set_up_parser() -> argparse.ArgumentParser:
         "--gpu", default=0, type=int, help="GPU ID to use", metavar="ID"
     )
     parser.add_argument(
-        "--root-dir",
+        "--root",
         required=True,
         type=str,
         help="root directory of the dataset for the accompanying task",
@@ -123,7 +123,7 @@ def main(args: argparse.Namespace) -> None:
         args: command-line arguments
     """
     assert os.path.exists(args.input_checkpoint)
-    assert os.path.exists(args.root_dir)
+    assert os.path.exists(args.root)
     TASK = TASK_TO_MODULES_MAPPING[args.task][0]
     DATAMODULE = TASK_TO_MODULES_MAPPING[args.task][1]
 
@@ -135,7 +135,7 @@ def main(args: argparse.Namespace) -> None:
 
     dm = DATAMODULE(  # type: ignore[call-arg]
         seed=args.seed,
-        root_dir=args.root_dir,
+        root=args.root,
         num_workers=args.num_workers,
         batch_size=args.batch_size,
     )

diff --git a/experiments/test_chesapeakecvpr_models.py b/experiments/test_chesapeakecvpr_models.py
@@ -138,7 +138,7 @@ def main(args: argparse.Namespace) -> None:
         for test_splits in ALL_TEST_SPLITS:
 
             dm = ChesapeakeCVPRDataModule(
-                args.chesapeakecvpr_root,
+                root=args.chesapeakecvpr_root,
                 train_splits=["de-train"],
                 val_splits=["de-val"],
                 test_splits=test_splits,

diff --git a/tests/conf/bigearthnet_all.yaml b/tests/conf/bigearthnet_all.yaml
@@ -9,7 +9,7 @@ experiment:
     in_channels: 14
     num_classes: 19
   datamodule:
-    root_dir: "tests/data/bigearthnet"
+    root: "tests/data/bigearthnet"
     bands: "all"
     num_classes: ${experiment.module.num_classes}
     batch_size: 1

diff --git a/tests/conf/bigearthnet_s1.yaml b/tests/conf/bigearthnet_s1.yaml
@@ -9,7 +9,7 @@ experiment:
     in_channels: 2
     num_classes: 19
   datamodule:
-    root_dir: "tests/data/bigearthnet"
+    root: "tests/data/bigearthnet"
     bands: "s1"
     num_classes: ${experiment.module.num_classes}
     batch_size: 1

diff --git a/tests/conf/bigearthnet_s2.yaml b/tests/conf/bigearthnet_s2.yaml
@@ -9,7 +9,7 @@ experiment:
     in_channels: 12
     num_classes: 19
   datamodule:
-    root_dir: "tests/data/bigearthnet"
+    root: "tests/data/bigearthnet"
     bands: "s2"
     num_classes: ${experiment.module.num_classes}
     batch_size: 1

diff --git a/tests/conf/byol.yaml b/tests/conf/byol.yaml
@@ -9,7 +9,7 @@ experiment:
     learning_rate: 1e-3
     learning_rate_schedule_patience: 6
   datamodule:
-    root_dir: "tests/data/chesapeake/cvpr"
+    root: "tests/data/chesapeake/cvpr"
     train_splits:
       - "de-test"
     val_splits:

diff --git a/tests/conf/chesapeake_cvpr_5.yaml b/tests/conf/chesapeake_cvpr_5.yaml
@@ -14,7 +14,7 @@ experiment:
     ignore_index: null
     imagenet_pretraining: False
   datamodule:
-    root_dir: "tests/data/chesapeake/cvpr"
+    root: "tests/data/chesapeake/cvpr"
     train_splits:
     - "de-test"
     val_splits:

diff --git a/tests/conf/chesapeake_cvpr_7.yaml b/tests/conf/chesapeake_cvpr_7.yaml
@@ -14,7 +14,7 @@ experiment:
     ignore_index: null
     imagenet_pretraining: True
   datamodule:
-    root_dir: "tests/data/chesapeake/cvpr"
+    root: "tests/data/chesapeake/cvpr"
     train_splits:
     - "de-test"
     val_splits:

diff --git a/tests/conf/chesapeake_cvpr_prior.yaml b/tests/conf/chesapeake_cvpr_prior.yaml
@@ -14,7 +14,7 @@ experiment:
     ignore_index: null
     imagenet_pretraining: False
   datamodule:
-    root_dir: "tests/data/chesapeake/cvpr"
+    root: "tests/data/chesapeake/cvpr"
     train_splits:
     - "de-test"
     val_splits:

diff --git a/tests/conf/cowc_counting.yaml b/tests/conf/cowc_counting.yaml
@@ -6,7 +6,7 @@ experiment:
     learning_rate_schedule_patience: 2
     pretrained: True
   datamodule:
-    root_dir: "tests/data/cowc_counting"
+    root: "tests/data/cowc_counting"
     seed: 0
     batch_size: 1
     num_workers: 0
diff --git a/tests/conf/cyclone.yaml b/tests/conf/cyclone.yaml
@@ -6,7 +6,7 @@ experiment:
     learning_rate_schedule_patience: 2
     pretrained: False
   datamodule:
-    root_dir: "tests/data/cyclone"
+    root: "tests/data/cyclone"
     seed: 0
     batch_size: 1
     num_workers: 0
diff --git a/tests/conf/deepglobelandcover_0.yaml b/tests/conf/deepglobelandcover_0.yaml
@@ -13,7 +13,7 @@ experiment:
     num_filters: 1
     ignore_index: null
   datamodule:
-    root_dir: "tests/data/deepglobelandcover"
+    root: "tests/data/deepglobelandcover"
     val_split_pct: 0.0
     batch_size: 1
     num_workers: 0
diff --git a/tests/conf/deepglobelandcover_5.yaml b/tests/conf/deepglobelandcover_5.yaml
@@ -13,7 +13,7 @@ experiment:
     num_filters: 1
     ignore_index: null
   datamodule:
-    root_dir: "tests/data/deepglobelandcover"
+    root: "tests/data/deepglobelandcover"
     val_split_pct: 0.5
     batch_size: 1
     num_workers: 0
diff --git a/tests/conf/etci2021.yaml b/tests/conf/etci2021.yaml
@@ -11,6 +11,6 @@ experiment:
     num_classes: 2
     ignore_index: 0
   datamodule:
-    root_dir: "tests/data/etci2021"
+    root: "tests/data/etci2021"
     batch_size: 1
     num_workers: 0
diff --git a/tests/conf/eurosat.yaml b/tests/conf/eurosat.yaml
@@ -9,6 +9,6 @@ experiment:
     in_channels: 13
     num_classes: 2
   datamodule:
-    root_dir: "tests/data/eurosat"
+    root: "tests/data/eurosat"
     batch_size: 1
     num_workers: 0
diff --git a/tests/conf/inria.yaml b/tests/conf/inria.yaml
@@ -11,7 +11,7 @@ experiment:
     num_classes: 2
     ignore_index: null
   datamodule:
-      root_dir: "tests/data/inria"
+      root: "tests/data/inria"
       batch_size: 1
       num_workers: 0
       val_split_pct: 0.2

diff --git a/tests/conf/landcoverai.yaml b/tests/conf/landcoverai.yaml
@@ -13,6 +13,6 @@ experiment:
     num_filters: 1
     ignore_index: null
   datamodule:
-    root_dir: "tests/data/landcoverai"
+    root: "tests/data/landcoverai"
     batch_size: 1
     num_workers: 0
diff --git a/tests/conf/naipchesapeake.yaml b/tests/conf/naipchesapeake.yaml
@@ -13,8 +13,8 @@ experiment:
     num_filters: 1
     ignore_index: null
   datamodule:
-    naip_root_dir: "tests/data/naip"
-    chesapeake_root_dir: "tests/data/chesapeake/BAYWIDE"
+    naip_root: "tests/data/naip"
+    chesapeake_root: "tests/data/chesapeake/BAYWIDE"
     batch_size: 2
     num_workers: 0
     patch_size: 32
diff --git a/tests/conf/oscd_all.yaml b/tests/conf/oscd_all.yaml
@@ -13,8 +13,8 @@ experiment:
     num_filters: 1
     ignore_index: null
   datamodule:
-    root_dir: "tests/data/oscd"
-    batch_size: 1
+    root: "tests/data/oscd"
+    train_batch_size: 1
     num_workers: 0
     val_split_pct: 0.5
     bands: "all"

diff --git a/tests/conf/oscd_rgb.yaml b/tests/conf/oscd_rgb.yaml
@@ -13,8 +13,8 @@ experiment:
     num_filters: 1
     ignore_index: null
   datamodule:
-    root_dir: "tests/data/oscd"
-    batch_size: 1
+    root: "tests/data/oscd"
+    train_batch_size: 1
     num_workers: 0
     val_split_pct: 0.5
     bands: "rgb"

diff --git a/tests/conf/resisc45.yaml b/tests/conf/resisc45.yaml
@@ -9,6 +9,6 @@ experiment:
     in_channels: 3
     num_classes: 3
   datamodule:
-    root_dir: "tests/data/resisc45"
+    root: "tests/data/resisc45"
     batch_size: 1
     num_workers: 0
diff --git a/tests/conf/sen12ms_all.yaml b/tests/conf/sen12ms_all.yaml
@@ -12,7 +12,7 @@ experiment:
     num_classes: 11
     ignore_index: null
   datamodule:
-    root_dir: "tests/data/sen12ms"
+    root: "tests/data/sen12ms"
     band_set: "all"
     batch_size: 1
     num_workers: 0

diff --git a/tests/conf/sen12ms_s1.yaml b/tests/conf/sen12ms_s1.yaml
@@ -13,7 +13,7 @@ experiment:
     num_classes: 11
     ignore_index: null
   datamodule:
-    root_dir: "tests/data/sen12ms"
+    root: "tests/data/sen12ms"
     band_set: "s1"
     batch_size: 1
     num_workers: 0

diff --git a/tests/conf/sen12ms_s2_all.yaml b/tests/conf/sen12ms_s2_all.yaml
@@ -12,7 +12,7 @@ experiment:
     num_classes: 11
     ignore_index: null
   datamodule:
-    root_dir: "tests/data/sen12ms"
+    root: "tests/data/sen12ms"
     band_set: "s2-all"
     batch_size: 1
     num_workers: 0

diff --git a/tests/conf/sen12ms_s2_reduced.yaml b/tests/conf/sen12ms_s2_reduced.yaml
@@ -12,7 +12,7 @@ experiment:
     num_classes: 11
     ignore_index: null
   datamodule:
-    root_dir: "tests/data/sen12ms"
+    root: "tests/data/sen12ms"
     band_set: "s2-reduced"
     batch_size: 1
     num_workers: 0