DataModules: run all data augmentation on the GPU (#992)

* DataModules: run all data augmentation on the GPU

* Passing tests

* Update BigEarthNet

* Break ChesapeakeCVPR

* Update COWC

* Update Cyclone

* Update ETCI2021

* mypy fixes

* Update FAIR1M

* Update Inria

* Update LandCoverAI

* Update LoveDA

* Update NAIP

* Update NASA

* Update OSCD

* Update RESISC45

* Update SEN12MS

* Update So2Sat

* Update SpaceNet

* Update UCMerced

* Update USAVars

* Update xview

* Remove seed

* mypy fixes

* OSCD hacks

* Add NonGeoDataModule base class

* Fixes

* Add base class to docs

* mypy fixes

* Fix several tests

* Fix Normalize

* Syntax error

* Fix bigearthnet

* Fix dtype

* Consistent kornia import

* Get regression datasets working

* Fix detection tests

* Fix some chesapeake bugs

* Fix several segmentation issues

* isort fixes

* Undo breaking change

* Remove more code duplication, standardize docstrings

* mypy fixes

* Add default augmentation

* Augmentations can be any callable

* Fix datasets tests

* Fix datamodule tests

* Fix more datamodules

* Typo fix

* Set up val_dataset even when fit

* Fix classification tests

* Fix ETCI2021

* Fix SEN12MS

* Add GeoDataModule base class

* Fix several chesapeake bugs

* Fix dtype and shape

* Fix crs/bbox issue

* Fix test dtype

* Fix unequal size stacking error

* flake8 fix

* Better checks on sampler

* Fix bug introduced in NAIP dm

* Fix chesapeake dimensions

* Add one to mask

* Fix missing imports

* Fix batch size

* Simplify augmentations

* Don't run test or predict without datasets

* Fix tests

* Allow shared dataset

* One more try

* Fix typo

* Fix another typo

* Fix Chesapeake dimensions

* Apply augmentations during sanity check too

* Don't reuse fixtures

* Increase coverage

* Fix ETCI tests

* Test predict_step

* Test all loss methods

* Simplify validation plotting

* Document new classes

* Fix plotting

* Plotting should be robust in case dataset does not contain RGB bands

* Fix flake8

* 100% coverage of trainers

* Add lightning-lite dependency

* Revert "Add lightning-lite dependency"

This reverts commit 1df7291.

* Define our own MisconfigurationException

* Properly test new data module base classes

* Fix mistake in setup call

* ExtractTensorPatches runs into OOM errors

* Test both fast_dev_run True and False

* Fix plot methods

* Fix OSCD tests

* Fix bug with inconsistent train/val/test splits between stages

* Fix issues with images of different sizes

* Fix OSCD tests

* Fix OSCD tests

* Bad rebase

* No trainer for OSCD so no need for config

* Bad rebase

* plot: only works during validation

* Fix collation of NASA Marine Debris dataset

* flake8 fix

* Quick test

* Revert "Quick test"

This reverts commit f465efc.

* 56 workers is a bit excessive

Co-authored-by: Caleb Robinson <calebrob6@gmail.com>

conf/chesapeake_cvpr.yaml

@@ -25,9 +25,8 @@ experiment:
       - "de-val"
     test_splits:
       - "de-test"
-    patches_per_tile: 200
+    batch_size: 200
     patch_size: 256
-    batch_size: 64
     num_workers: 4
     class_set: ${experiment.module.num_classes}
     use_prior_labels: False

conf/cowc_counting.yaml

@@ -13,6 +13,5 @@ experiment:
     learning_rate_schedule_patience: 2
   datamodule:
     root: "data/cowc_counting"
-    seed: 0
     batch_size: 64
     num_workers: 4

conf/cyclone.yaml

@@ -13,6 +13,5 @@ experiment:
     learning_rate_schedule_patience: 2
   datamodule:
     root: "data/cyclone"
-    seed: 0
     batch_size: 32
     num_workers: 4

conf/deepglobelandcover.yaml

@@ -14,8 +14,7 @@ experiment:
     ignore_index: null
   datamodule:
     root: "data/deepglobelandcover"
-    num_tiles_per_batch: 16
-    num_patches_per_tile: 16
+    batch_size: 1
     patch_size: 64
     val_split_pct: 0.5
     num_workers: 0

conf/gid15.yaml

@@ -14,8 +14,7 @@ experiment:
     ignore_index: null
   datamodule:
     root: "data/gid15"
-    num_tiles_per_batch: 16
-    num_patches_per_tile: 16
+    batch_size: 1
     patch_size: 64
     val_split_pct: 0.5
     num_workers: 0

conf/inria.yaml

@@ -1,7 +1,7 @@
 program:
   overwrite: True
 
-  
+
 trainer:
   gpus: 1
   min_epochs: 5
@@ -23,8 +23,7 @@ experiment:
     num_classes: 2
     ignore_index: null
   datamodule:
-      root: "data/inria"
-      batch_size: 2
-      num_workers: 32
-      patch_size: 512
-      num_patches_per_tile: 4
+    root: "data/inria"
+    batch_size: 1
+    patch_size: 512
+    num_workers: 32

conf/nasa_marine_debris.yaml

@@ -1,5 +1,4 @@
 program:
-  seed: 0
   overwrite: True
 
 trainer:
@@ -22,5 +21,5 @@ experiment:
   datamodule:
     root: "data/nasamr/nasa_marine_debris"
     batch_size: 4
-    num_workers: 56
+    num_workers: 6
     val_split_pct: 0.2

conf/potsdam2d.yaml

@@ -14,8 +14,7 @@ experiment:
     ignore_index: null
   datamodule:
     root: "data/potsdam"
-    num_tiles_per_batch: 16
-    num_patches_per_tile: 16
+    batch_size: 1
     patch_size: 64
     val_split_pct: 0.5
     num_workers: 0

conf/sen12ms.yaml

@@ -20,4 +20,3 @@ experiment:
     band_set: "all"
     batch_size: 32
     num_workers: 4
-    seed: 0

conf/so2sat.yaml

@@ -11,11 +11,10 @@ experiment:
     learning_rate: 1e-3
     learning_rate_schedule_patience: 6
     weights: null
-    in_channels: 3
+    in_channels: 18
     num_classes: 17
   datamodule:
     root: "data/so2sat"
     batch_size: 128
     num_workers: 4
-    band_set: "rgb"
-    unsupervised_mode: False
+    band_set: "all"

conf/spacenet1.yaml

@@ -1,6 +1,5 @@
 program:
   overwrite: False
-  seed: 0
 trainer:
   gpus: [3]
   min_epochs: 50
@@ -22,4 +21,4 @@ experiment:
   datamodule:
     root: "data/spacenet"
     batch_size: 32
-    num_workers: 4
+    num_workers: 4

conf/vaihingen2d.yaml

@@ -14,8 +14,7 @@ experiment:
     ignore_index: null
   datamodule:
     root: "data/vaihingen"
-    num_tiles_per_batch: 16
-    num_patches_per_tile: 16
+    batch_size: 1
     patch_size: 64
     val_split_pct: 0.5
     num_workers: 0

docs/api/datamodules.rst

@@ -128,3 +128,21 @@ xView2
 ^^^^^^
 
 .. autoclass:: XView2DataModule
+
+Base Classes
+------------
+
+GeoDataModule
+^^^^^^^^^^^^^
+
+.. autoclass:: GeoDataModule
+
+NonGeoDataModule
+^^^^^^^^^^^^^^^^
+
+.. autoclass:: NonGeoDataModule
+
+Utilities
+---------
+
+.. autoclass:: MisconfigurationException

tests/conf/chesapeake_cvpr_5.yaml

@@ -20,9 +20,8 @@ experiment:
     - "de-test"
     test_splits:
     - "de-test"
-    patches_per_tile: 2
-    patch_size: 64
     batch_size: 2
+    patch_size: 64
     num_workers: 0
     class_set: ${experiment.module.num_classes}
     use_prior_labels: False

tests/conf/chesapeake_cvpr_7.yaml

@@ -20,9 +20,8 @@ experiment:
     - "de-test"
     test_splits:
     - "de-test"
-    patches_per_tile: 2
-    patch_size: 64
     batch_size: 2
+    patch_size: 64
     num_workers: 0
     class_set: ${experiment.module.num_classes}
     use_prior_labels: False

tests/conf/chesapeake_cvpr_prior.yaml

@@ -20,9 +20,8 @@ experiment:
     - "de-test"
     test_splits:
     - "de-test"
-    patches_per_tile: 2
-    patch_size: 64
     batch_size: 2
+    patch_size: 64
     num_workers: 0
     class_set: ${experiment.module.num_classes}
     use_prior_labels: True

tests/conf/cowc_counting.yaml

@@ -10,6 +10,5 @@ experiment:
   datamodule:
     root: "tests/data/cowc_counting"
     download: true
-    seed: 0
     batch_size: 1
     num_workers: 0

tests/conf/cyclone.yaml

@@ -10,6 +10,5 @@ experiment:
   datamodule:
     root: "tests/data/cyclone"
     download: true
-    seed: 0
     batch_size: 1
     num_workers: 0

tests/conf/deepglobelandcover.yaml

@@ -14,8 +14,7 @@ experiment:
     ignore_index: null
   datamodule:
     root: "tests/data/deepglobelandcover"
-    num_tiles_per_batch: 1
-    num_patches_per_tile: 1
+    batch_size: 1
     patch_size: 2
     val_split_pct: 0.5
     num_workers: 0

tests/conf/gid15.yaml

@@ -15,8 +15,7 @@ experiment:
   datamodule:
     root: "tests/data/gid15"
     download: true
-    num_tiles_per_batch: 1
-    num_patches_per_tile: 1
+    batch_size: 1
     patch_size: 2
     val_split_pct: 0.5
     num_workers: 0

tests/conf/inria_test.yaml → tests/conf/inria.yaml

@@ -11,10 +11,9 @@ experiment:
     num_classes: 2
     ignore_index: null
   datamodule:
-      root: "tests/data/inria"
-      batch_size: 1
-      num_workers: 0
-      val_split_pct: 0.2
-      test_split_pct: 0.2
-      patch_size: 2
-      num_patches_per_tile: 2
+    root: "tests/data/inria"
+    batch_size: 1
+    patch_size: 2
+    num_workers: 0
+    val_split_pct: 0.2
+    test_split_pct: 0.2

tests/conf/potsdam2d.yaml

@@ -14,8 +14,7 @@ experiment:
     ignore_index: null
   datamodule:
     root: "tests/data/potsdam"
-    num_tiles_per_batch: 1
-    num_patches_per_tile: 1
+    batch_size: 1
     patch_size: 2
     val_split_pct: 0.5
     num_workers: 0

tests/conf/sen12ms_all.yaml

@@ -15,4 +15,3 @@ experiment:
     band_set: "all"
     batch_size: 1
     num_workers: 0
-    seed: 0

tests/conf/sen12ms_s1.yaml

@@ -16,4 +16,3 @@ experiment:
     band_set: "s1"
     batch_size: 1
     num_workers: 0
-    seed: 0

tests/conf/sen12ms_s2_all.yaml

@@ -15,4 +15,3 @@ experiment:
     band_set: "s2-all"
     batch_size: 1
     num_workers: 0
-    seed: 0

tests/conf/sen12ms_s2_reduced.yaml

@@ -15,4 +15,3 @@ experiment:
     band_set: "s2-reduced"
     batch_size: 1
     num_workers: 0
-    seed: 0

tests/conf/so2sat_all.yaml

@@ -0,0 +1,15 @@
+experiment:
+  task: "so2sat"
+  module:
+    loss: "ce"
+    model: "resnet18"
+    learning_rate: 1e-3
+    learning_rate_schedule_patience: 6
+    weights: null
+    in_channels: 18
+    num_classes: 17
+  datamodule:
+    root: "tests/data/so2sat"
+    batch_size: 1
+    num_workers: 0
+    band_set: "all"

tests/conf/so2sat_supervised.yaml → tests/conf/so2sat_s1.yaml

@@ -6,11 +6,10 @@ experiment:
     learning_rate: 1e-3
     learning_rate_schedule_patience: 6
     weights: null
-    in_channels: 3
+    in_channels: 8
     num_classes: 17
   datamodule:
     root: "tests/data/so2sat"
     batch_size: 1
     num_workers: 0
-    band_set: "rgb"
-    unsupervised_mode: False
+    band_set: "s1"

tests/conf/so2sat_unsupervised.yaml → tests/conf/so2sat_s2.yaml

@@ -6,11 +6,10 @@ experiment:
     learning_rate: 1e-3
     learning_rate_schedule_patience: 6
     weights: null
-    in_channels: 3
+    in_channels: 10
     num_classes: 17
   datamodule:
     root: "tests/data/so2sat"
     batch_size: 1
     num_workers: 0
-    band_set: "rgb"
-    unsupervised_mode: True
+    band_set: "s2"

tests/conf/vaihingen2d.yaml

@@ -14,8 +14,7 @@ experiment:
     ignore_index: null
   datamodule:
     root: "tests/data/vaihingen"
-    num_tiles_per_batch: 1
-    num_patches_per_tile: 1
+    batch_size: 1
     patch_size: 2
     val_split_pct: 0.5
     num_workers: 0