Trainers: predict step

[adamjstewart]

Summary

This issue is to track progress on adding a predict step to all Trainers.

• BYOLTask (add predict_step to BYOLTask #819)
• ClassificationTask (ClassificationTask predict step #790)
• MultiLabelClassificationTask (MultiLabelClassificationTask predict step #792)
• ObjectDetectionTask (Add support for ObjectDetection #758)
• RegressionTask (add predict_step to RegressionTask #818)
• SemanticSegmentationTask (add predict_step to 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

[isaaccorley]

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.

[Geethen]

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.

[adamjstewart]

What's left to do for the SemanticSegmentationTask? Would like to finish this before the next release.

[isaaccorley]

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.

[adamjstewart]

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.

[isaaccorley]

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.