DataModules: run all data augmentation on the GPU

[adamjstewart]

This PR overhauls all of our data modules to improve uniformity. This includes the following changes:

• Add GeoDataModule and NonGeoDataModule base classes to reduce code duplication
• Only instantiate the datasets that are needed for a particular stage
• Replace torchvision with kornia (better support for MSI, GPU, inverse)
• Replace dataset transforms with on_after_batch_transfer (CPU -> GPU, sample -> batch, faster)
• Remove instance methods for preprocessing (fixes Trainers: num_workers > 0 results in pickling error on macOS/Windows #886)
• Fix bug where train/val/test split would differ for each stage and images would leak between sets
• Deprecate torchgeo.transforms.AugmentationSequential (use kornia.augmentation.AugmentationSequential instead)
• Deprecate torchgeo.datamodules.utils.dataset_split (use torch.utils.data.random_split instead)

In a future PR, I'm planning on extending this to the rest of our transforms:

• Rewrite all index transforms to be compatible with Kornia (Convert all index transforms to Kornia #999)
• Update tutorials to use Kornia with our transforms
• Upstream and remove our custom transforms and AugmentationSequential hacks

Fixes #619
Fixes #337
Fixes #336

[adamjstewart]

If we enforce that all datasets must have a plot method, we could remove this check. Currently the only ones lacking are VHR-10 (WIP) and our point datasets (GBIF, iNaturalist, EDDMapS).

[adamjstewart]

The idea for all of these is that you can either define a single attribute (self.foo) or a different attribute for each stage (self.train_foo, self.val_foo, etc.)

[adamjstewart]

Need to ignore type warnings because not all datasets accept a split argument

[adamjstewart]

Note that datasets/samplers with length 0 also evaluate to False, so this may lead to red herrings

[adamjstewart]

I may decide to remove these completely but didn't do that for this PR since it requires a ton of unrelated changes. Will reconsider when working on #985.

[adamjstewart]

I'm having trouble understanding the failing BYOL tests. The same datamodules work fine with SemanticSegmentationTask. If anyone can figure out how to fix these, we can see how much coverage we lack and fix that.

[adamjstewart]

The last remaining test failure is due to discrepancies between Chesapeake and all other datasets. All of our datasets return a mask of shape "h w" but Chesapeake returns "c h w". This breaks everything. Can't wait until PEP 646 is supported in PyTorch...

How should we handle this? It doesn't seem like Chesapeake can be changed to "h w" since the prior labels are c=4. We could change every other dataset to be "c h w" like Kornia expects, but that sounds like a lot more work. We could also write a custom AugmentationSequential just for Chesapeake that handles things properly. Our AugmentationSequential only works for batches, not samples.

[adamjstewart]

@ashnair1 @nilsleh how would you propose to solve the ExtractTensorPatches OOM issue?

Easiest fix is to use RandomNCrop instead (this is what we do during train). I'm happy to do this. However, it doesn't really make sense during val/test, and isn't useful at all during predict.

Alternative fix would be to introduce a GridNonGeoSampler that works similarly to GridGeoSampler and allows a single image to span multiple batches. However, this would require all NonGeoDatasets to change from __getitem__(idx: int) to __getitem__(idx: int, rows: slice, cols: slice) which would be a ton of work.

Any other ideas for how to handle this?

[ashnair1]

@ashnair1 @nilsleh how would you propose to solve the ExtractTensorPatches OOM issue?

Easiest fix is to use RandomNCrop instead (this is what we do during train). I'm happy to do this. However, it doesn't really make sense during val/test, and isn't useful at all during predict.

Alternative fix would be to introduce a GridNonGeoSampler that works similarly to GridGeoSampler and allows a single image to span multiple batches. However, this would require all NonGeoDatasets to change from __getitem__(idx: int) to __getitem__(idx: int, rows: slice, cols: slice) which would be a ton of work.

Any other ideas for how to handle this?

A trivial solution to OOM would be to run inference in an iterative manner whenever we use _ExtractTensorPatches.

For example, in the validation_step of torchgeo/trainers/segmentation.py we iterate over the batch instead of passing in the entire batch.

def validation_step(self, *args: Any, **kwargs: Any) -> None:
    """Compute validation loss and log example predictions.

    Args:
        batch: the output of your DataLoader
        batch_idx: the index of this batch
    """
    batch = args[0]
    batch_idx = args[1]
    x = batch["image"]
    y = batch["mask"]
-   y_hat = self(x)
+   from einops import rearrange
+   y_hat_list = []
+   for i in range(x.shape[0]):
+       out = rearrange(x[i], 'c h w -> 1 c h w') 
+       out = self(out)
+       y_hat_list.append(out)    
+   y_hat = rearrange(y_hat_list, 'b 1 c h w -> b c h w')
    y_hat_hard = y_hat.argmax(dim=1)

    loss = self.loss(y_hat, y)

[adamjstewart]

run inference in an iterative manner whenever we use _ExtractTensorPatches

I don't think there's an easy way to tell whether or not we are using _ExtractTensorPatches. We could say "if batch size > X run iteratively" but I'm not sure what X should be.

[ashnair1]

How about making this the default behaviour for val, test and predict steps?

X is dependent on the hardware RAM which can only determined by the user. So I don't think we can go down that route.

[adamjstewart]

I think replacing all val/test/predict steps from batch processing to a for-loop over the batch size would be extremely bad for speed.

[nilsleh]

We could say "if batch size > X run iteratively" but I'm not sure what X should be.

Don't know if it is a desired approach, but could you wrap y_hat = self(x) in a try, except block checking for OOM, and in the except block do the looping?

[adamjstewart]

If I recall, it isn't possible to check for OOM, the program crashes without raising an error. But correct me if I'm wrong. Also, it would be difficult to get test coverage of that branch.

[nilsleh]

On the pytorch docs, they seem to suggest a possible way as a RuntimeError but as you said difficult to get test coverage.

[ashnair1]

Currently training fails during the sanity checking (validation) step fig = datamodule.plot(sample) as fig will be of type None because self.val_dataset (of type torch.utils.data.Subset) does not have a plot method.

[adamjstewart]

@calebrob6 starting a thread here so we can discuss how to handle NASA Marine Debris.

I think the solution in this collate_fn is actually correct, we want a list of tensors, not a single tensor: kornia/kornia#1497

Note that the name is wrong, we should be used bbox_xyxy or bbox_xywh instead of bbox (boxes is translated to bbox internally in our AugmentationSequential wrapper) but we can fix that when working on #985.

[adamjstewart]

Okay, see if the last commit fixes this.

[adamjstewart]

@calebrob6 I believe I've addressed all review comments, let me know if you find anything else to fix.

[calebrob6]

Yep that did it, nice job!