datamodules: Switch to kornia AugmentationSequential

[ashnair1]

Switch to kornia's AugmentationSequential in datamodules submodule

Closes #1432

TODO (once kornia 0.7.4 is released):

• Remove manual setting of keepdim attribute
• Bump min version of kornia

[ashnair1]

TODO: Tutorials need to be updated once we fully switch to kornia's AugmentationSequential

[adamjstewart]

Still confused by some of the dataset changes, but all datamodule changes now look good. I'll see if I can figure out why the tests are failing.

[adamjstewart]

Bit confused by this. Relative time and ocean are not returned if transforms is None? This PR only touches datamodule augmentations, not dataset transforms, right? Why is this needed? Same question for all other datasets.

[ashnair1]

Reason for this mentioned here- #2147 (comment)

These type of fixes won't be required once kornia/kornia#2971 makes it to a release. It allows non-kornia compatible keys to be passed.

[adamjstewart]

Hmm, this kind of seems like a hack. It's an implementation detail that TropicalCycloneDataModule doesn't use TropicalCyclone(..., transform=...), not by design. I definitely don't think we should return different things depending on whether transforms is used or not.

I have a better idea. In BaseDataModule in transfer_batch_to_device, drop all keys that Kornia doesn't recognize. This will allow us to drop most changes to datasets and the transfer_batch_to_device in GeoDataModule. You can loop over all keys and drop any that don't correspond to a valid Kornia key (there should be a way to get a list of valid keys from Kornia).

[adamjstewart]

Does the mask occasionally change dimensions? Maybe this can be made uniform in getitem?

[adamjstewart]

But tests are failing for most GeoDataModules.

I wish it were this simple... Note that some failing tests are for NonGeoDataModules (ssl4eo_l_benchmark_*).

Also note that some tests fail depending on batch size. For example, gid15 passes with batch_size=1, but fails if batch_size=2. But there are other datasets for which both 1 and 2+ pass. For the case of gid15, it seems like this is because the image is 1x1, so this may be a red herring.

Most of the failing tests seem to be because the ground truth mask has a channel dimension, which cross-entropy loss does not like. Let me see if I can fix a few.

[adamjstewart]

The ChesapeakeCVPR dataset is royally screwed up and I'm not touching it, so I'll let you or @calebrob6 solve that one. The problem is that it doesn't use RasterDataset and the mask is sometimes single channel, sometimes multichannel, so you have to be very careful with all dimensions. I'm not even sure our trainer supports multichannel masks.

[adamjstewart]

Is this waiting on Kornia 0.7.4? Trying to decide if this should be in TorchGeo 0.6.0 next week or if we should bump it to a later release.

[ashnair1]

Is this waiting on Kornia 0.7.4? Trying to decide if this should be in TorchGeo 0.6.0 next week or if we should bump it to a later release.

Later release for sure. This does depend on Kornia 0.7.4.