Remove dataset-specific trainers

[adamjstewart]

Closes #205
Depends on #329

References:

• How to access `LightningDataModule` in `LightningModule` Lightning-AI/pytorch-lightning#10492
• LightningDataModule: GPU data augmentation support Lightning-AI/pytorch-lightning#10469

[isaaccorley]

Just a note: I've recently been using kornia augmentations in the datamodule on a side project and one of the things I ran into was how to let the datamodule know if I'm loading data for a train/val/test set so that I can choose to augment or not. Found that you can access a bool attr self.trainer.training in the datamodule so you can do something like:

def on_after_batch_transfer(self, batch, dataloader_idx):
   if self.trainer.training:
      # Augment only if loading for train_step
      batch = augmentations(batch)
    return batch

[adamjstewart]

Note that none of this code currently gets hit by our tests. We aren't using a pl.Trainer and so things like self.trainer are None. Still trying to figure out the best way to test this.

[adamjstewart]

Note: I don't think we're adding the predictions to the batch before plotting, we probably should

[adamjstewart]

Forgot to remove these in #329

[adamjstewart]

Test failure is because the So2Sat dataset doesn't know how to plot any of the datamodule reduced band set options. We should probably move these to the dataset level.

[adamjstewart]

Another hiccup: self.trainer.datamodule.val_dataset.plot(...) doesn't work for datamodules that use Subset or random_split. One possible solution would be to use self.trainer.datamodule.plot(...) and add a plot(...) method to every DataModule that passes all args to the Dataset plot(...) method.

[isaaccorley]

You can access the plot method for Subset datasets like self.trainer.datamodule.val_dataset.dataset.plot. Not sure what workaround we should make for this.

Edit: I think adding a plot method to each datamodule that just calls the dataset plot method is a decent solution.

[adamjstewart]

I believe the failing unit tests for ClassificationTask are because VisionClassificationDataset is overwriting self.classes and our fake data only has 2 classes. Simple fix would be to add more fake data.

Still haven't investigated the failing unit tests for SemanticSegmentationTask. Will do so tomorrow.