Repository: microsoft/torchgeo

Issue:
Trainers: predict step
### Summary

This issue is to track progress on adding a `predict` step to all Trainers.

- [x] BYOLTask (#819)
- [x] ClassificationTask (#790)
- [x] MultiLabelClassificationTask (#792)
- [x] ObjectDetectionTask (#758)
- [x] RegressionTask (#818)
- [x] SemanticSegmentationTask (#939)

### Rationale

The default `predict` step does not know how to handle our batch dicts.

### Implementation

See implementations that have already been finished.

### Alternatives

_No response_

### Additional information

_No response_


Discussion:
BYOL could just return embeddings for each sample. RegressionTask should also be straightforward similar to ClassificationTask. I'll tackle these today. I think SemanticSegmentation may need some thought since it can be a large amount of data per sample and we don't want users to run OOM by accumulating masks for each sample.
If you have not already considered a similar option for the semantic segmentation regression task: https://rasterio.readthedocs.io/en/latest/topics/concurrency.html

uses rasterio + windows and concurrent to batch write data to a single large image.
What's left to do for the SemanticSegmentationTask? Would like to finish this before the next release.
Adding a predict step similar to the other tasks will just accumulate the segmentation output masks into a list. For large images and datasets this may not be desirable for memory usage reasons. We may want to consider saving the outputs somewhere instead or just leaving it to the user to override how they want the model to perform predictions.
Could we go with a simpler implementation in which we return the predictions without saving? Or assume the user will want to predict on a small enough area that it can fit in memory? The current state where it crashes is unideal. 
Yes I'll make a PR for this. I'll add a note to the docstring for users to override by making a custom task if they want to do something specific with the masks.

Resolving patch:
diff --git a/torchgeo/trainers/segmentation.py b/torchgeo/trainers/segmentation.py
--- a/torchgeo/trainers/segmentation.py
+++ b/torchgeo/trainers/segmentation.py
@@ -248,6 +248,25 @@ def test_epoch_end(self, outputs: Any) -> None:
         self.log_dict(self.test_metrics.compute())
         self.test_metrics.reset()
 
+    def predict_step(self, *args: Any, **kwargs: Any) -> Tensor:
+        """Compute and return the predictions.
+
+        By default, this will loop over images in a dataloader and aggregate
+        predictions into a list. This may not be desirable if you have many images
+        or large images which could cause out of memory errors. In this case
+        it's recommended to override this with a custom predict_step.
+
+        Args:
+            batch: the output of your DataLoader
+
+        Returns:
+            predicted softmax probabilities
+        """
+        batch = args[0]
+        x = batch["image"]
+        y_hat: Tensor = self(x).softmax(dim=1)
+        return y_hat
+
     def configure_optimizers(self) -> Dict[str, Any]:
         """Initialize the optimizer and learning rate scheduler.
 


Test patch:
diff --git a/tests/trainers/test_segmentation.py b/tests/trainers/test_segmentation.py
--- a/tests/trainers/test_segmentation.py
+++ b/tests/trainers/test_segmentation.py
@@ -78,6 +78,7 @@ def test_trainer(
         trainer = Trainer(fast_dev_run=True, log_every_n_steps=1, max_epochs=1)
         trainer.fit(model=model, datamodule=datamodule)
         trainer.test(model=model, datamodule=datamodule)
+        trainer.predict(model=model, dataloaders=datamodule.val_dataloader())
 
     def test_no_logger(self) -> None:
         conf = OmegaConf.load(os.path.join("tests", "conf", "landcoverai.yaml"))