In this repository, you can find an implementation of Steerable Conditional Neural Processes, a work published at ICML 2021 that you find here.
Steerable Conditional Neural Processes (SteerCNPs) are an extension of Conditional Neural Processes. The model consists of an encoder which is based on the work of Gordon et al (2020) and a decoder which is an equivariant neural network of the form outlined in the work of Weiler et al (2019).
The image below depicts example predictions of the SteerCNP in the case of vector fields. Inputs are the red arrows and the model extracts the whole vector field.
This library provides an implementation of SteerCNPs and the code for two experiments.
We tested our model on two data sets: a Gaussian process regression task and real-world weather data. Below, the model predicts the wind in a cyclic region of 500km radius around Memphis in the South of the US. It gets measurements of wind, temperature and pressure from places marked in red.
To use this repository, one can simply clone it and install requirements specified in requirements.txt.
Please note this installs additional packages used for visualization. In particular, we use:
- PyTorch as a library for automatic differentation.
- the library E(2)-Steerable CNNs for any group-related tasks and an implementation of Steerable CNNs.
- the ERA5 data set giving grided global weather data. We share a full download and pre-processing script to extract the data set we use.
The core implementation of SteerCNPs are all files in the root. The folder "tasks" gives the two main tasks (data sets+ data loading scripts) which we have given our model: GP vector field data and real-world weather data. The folder "experiments" gives the main execution file per task. The folder CNP gives an implementation of Conditional Neural Processes to compare our results.