alan-turing-institute / alan-turing-institute/deepsensor

Patchwise training and inference

Open
#22 2 comments 0 reactions 0 assignees View on GitHub
enhancement thoughts welcome
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.