alan-turing-institute / alan-turing-institute/deepsensor
Patchwise training and inference
- Dominant language
- Python
- Stars
- 128
- Forks
- 28
- PR merge metrics
- No merged PRs in 30d
Description
Some `deepsensor` users may have dense environmental data spanning large spatial areas. For example:
* Problems with data spanning the whole globe
* Problems with high-resolution satellite data
In such cases, training and inference with a `ConvNP` over the entire region of data may be computationally prohibitive. Currently, the `TaskLoader` will sample context and target data over the entire spatial region that data is available, which could produce OOM issues. So we need to support chopping the data into smaller spatial patches.
**Training**
Supporting patchwise `ConvNP` training should just be a matter of updating the `TaskLoader` to slice the context and target datasets spatially to subsetted squares/regions before proceeding with the `TaskLoader.__call__` sampling functionality for generating `Task` objects. I believe this should be quite simple: for `xarray` data this would be `ds.sel(x1=slice(...), x2=slice(...))`, while for `pandas` data it would be `df.loc[slice(...), slice(...)]`.
**Inference**
Inference using the high-level `DeepSensorModel.predict` interface also needs support for patching. This requires functionality to stitch all the individual patch predictions together.
For on-grid `xarray` prediction, one solution might be to call `.predict` recursively over all the patches and then concatenate the resulting `xr.Dataset`s into single objects. This would require some kind of `patchify` bool to control this and avoid infinite recursion within the inner call. Open to other ideas!
However, model predictions could differ substantially from one side of a patch border to another (due to differing context information in each patch). We therefore may need to think about having overlapping patches and averaging model predictions somehow.
**Patch size/location question**
An open question is how the size and location of the patches should be determined. One option is to have the user pass the patch size in `TaskLoader.__call__` or `DeepSensorModel.predict`, and then the location will be generated randomly unless further kwargs are passed to override this and specify exact x1/x2 spatial bounds.
Contributor guide
Research direction
Start by reading TaskLoader.__call__ and DeepSensorModel.predict, then trace how xarray and pandas context and target datasets are sampled. The work is complete when training can operate on spatial patches and inference can combine the individual patch predictions into the expected output, with the patch-size, location, and border behavior decisions resolved.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- pandas, python
- Domain
- data, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100