- Resuming training states is also simplified, and compatible with multi-device setups:
>>> trainer = Trainer(... device=device, use_distributed=True)
>>> trainer.train(..., resume_from_dir="./checkpoints", ...)
Trainer resuming from epoch 5...
The Trainer accepts a `DataProcessor` module to streamline data pre- and post-processing. Explicit `train()` and `eval()` modes were also added to the `DataProcessor`, making the behavior during training and evaluation clearer and more explicit.
Last, logging has been reintroduced into the `Trainer`, including better integration with Weights and Biases (WandB) for experiment tracking.
For an example of the new Trainer in action, check out our [interactive training example](https://neuraloperator.github.io/dev/auto_examples/plot_FNO_darcy.html)!
New neural operator architectures
GINO
The GINO ([Geometry-Informed Neural Operator](https://arxiv.org/abs/2309.00583)) architecture brings the FNO's ability to learn complex dynamics to irregular geometries. The GINO maps between functions provided at arbitrary input coordinates. For maximum modularity, the GINO is composed of configurable `GNOBlock` and `FNOBlock` operator layers.
To see example usage, check out the training recipe in `scripts/train_gino_carcfd.py`.
UQNO
The UQNO implements [Uncertainty Quantification for Operator Learning](https://openreview.net/pdf?id=cGpegxy12T), thanks to ziqi-ma and dhpitt. For an example training recipe see `scripts/train_uqno_darcy.py`.
LocalFNO
The [Local Fourier Neural Operator](https://arxiv.org/abs/2402.16845) shares its forward pass and architecture with the standard FNO, with the key difference that its Fourier convolution layers are replaced with LocalFNOBlocks that place differential kernel layers and local integral layers in parallel to its Fourier layers.Thanks to mliuschi.
New neural operator blocks
Attention kernel integral
The `AttentionKernelIntegral`, by zijieli-Jlee brings the multi-head attention mechanism to operator learning.
Codomain-Attention Blocks
`CODABlocks`, by ashiq24, implements [Codomain Attention Neural Operators](https://arxiv.org/pdf/2403.12553), an operator block that extends transformer positional encoding, self-attention and normalization functionalities to function spaces.
Differential convolution
Implements the finite difference convolution required for [Local Neural Operators](https://arxiv.org/abs/2402.16845). The `DifferentialConv` computes a finite difference convolution on a regular grid, which converges to a directional derivative as the grid is refined. Thanks to mliuschi.
DISCO convolutions
`DiscreteContinuousConv2d` implements Discrete-Continuous Convolutions required for Local Neural Operators. Check our documentation for an [interactive demo](https://neuraloperator.github.io/dev/auto_examples/plot_DISCO_convolutions.html) of DISCO convolutions in use! Thanks to mliuschi and bonevbs.
GNOBlock
This layer implements the [Graph Neural Operator](https://arxiv.org/abs/2003.03485) architecture, which combines a spatial neighbor search with a pointwise aggregation to create a kernel integral similar to message-passing neural network.
LocalFNOBlock
- Adds a `LocalFNOBlocks` layer, which mirrors the `FNOBlock` architecture with differential and discrete-continuous convolutions in parallel to normal spectral convolutions ([468](https://github.com/neuraloperator/neuraloperator/pull/468))
Updates to existing models/layers
FNO
We've simplifies the FNO's documentation and initialization. Parameters are now ordered by importance, and init only requires values for the most crucial parameters
- Creating an FNO is now as simple as:
>>> from neuralop.models import FNO
>>> model = FNO(n_modes=(12,12), in_channels=1, out_channels=1, hidden_channels=64)
>>> model
FNO(
(positional_embedding): GridEmbeddingND()
(fno_blocks): FNOBlocks(
(convs): SpectralConv(
(weight): ModuleList(
(0-3): 4 x DenseTensor(shape=torch.Size([64, 64, 12, 7]), rank=None)
)
)
... torch.nn.Module printout truncated ...
We've also added support for functions that take on complex values in the spatial domain ([400](https://github.com/neuraloperator/neuraloperator/pull/400))
FNO-GNO
The FNOGNO combines `FNOBlocks` over a regular grid with a `GNOBlock` layer to map to arbitrary query points in the spatial domain. Updates include simplifications of parameters and documentation, and modularization to integrate the new `GNOBlock` layer.
New meta algorithms
Incremental learning of resolution and frequency
Implements [Incremental Spectral Learning](https://arxiv.org/pdf/2211.15188) to learn the smallest possible model for a given problem.
- Adds `IncrementalDataProcessor` and accompanying meta-alg ([274](https://github.com/neuraloperator/neuraloperator/pull/274))
- Meta-algorithm implemented as `IncrementalFNOTrainer`, a subclass of the `Trainer`
Thanks to Robertboy18
Out-of-the-box mixed-precision training
The Trainer's `mixed_precision` parameter automatically handles mixed-precision using the new `torch.autocast` framework.
A new data module
Overview
We've refactored all data-associated functionality into `neuralop.data`. Datasets are now located at `neuralop.data.datasets` and transforms are located at `neuralop.data.transforms`. We've also updated the interfaces of our example datasets, including the option to download data from the new [NeuralOperator Zenodo Community](https://zenodo.org/communities/neuraloperator).
The Darcy-flow, Burger's, and Navier-Stokes datasets now all derive from a `PTDataset` template class, which updates all interfaces, adds a`download` option and connects to the Zenodo data source ([299](https://github.com/neuraloperator/neuraloperator/pull/299)).
Car-CFD
To showcase our geometric models, we've also added a dataset of simulations of airflow over 3d ShapeNet car models ([452](https://github.com/neuraloperator/neuraloperator/pull/452)), as well as examples for the FNOGNO and GINO models (`scripts/train_fnogno_carcfd.py`, `scripts/train_gino_carcfd.py`).
Testing and Documentation
This release, we expanded test coverage to 80%. To ensure users can get the most out of the library with minimal effort, we also significantly improved the documentation, simplifying the docstring, adding more user-guides and new examples.
The **Quickstart Guide** has been fully updated to provide a smooth, hands-on experience for new users, with clear instructions on training, saving, and loading models.
Other changes
Correct complex optimization
PyTorch's default `Adam` implementation currently handles complex parameters by viewing as a 2-tensor stack `(real_values, imaginary_values)`. This leads to an incorrect magnitude computation, which in turn affects the computation of momentum. Our implementation keeps parameters as real, and updates the `SpectralConv`'s choice of parameter tensors to ensure that parameters are registered as complex.
- Adds custom `AdamW` implementation ([420](https://github.com/neuraloperator/neuraloperator/pull/420))
- Ensures `SpectralConv` parameters are registered with complex values ([401](https://github.com/neuraloperator/neuraloperator/pull/401))
Misc. bug fixes
- Fixes numerical instability in `UnitGaussianNormalizer` ([330](https://github.com/neuraloperator/neuraloperator/pull/330))
- Fixes factorized tensor contraction in `SpectralConv` when `separable=True` ([389](https://github.com/neuraloperator/neuraloperator/pull/389))
- Fixed a bug in `MGPatchingDataProcessor.to(device)` that prevented proper device transfers. ([449](https://github.com/neuraloperator/neuraloperator/pull/449))
- Corrected the number of training examples printed out during training in the `Trainer`. ([380](https://github.com/neuraloperator/neuraloperator/pull/380))
- Addressed future deprecation for `torch.amp.autocast` by updating to the latest initialization. ([393](https://github.com/neuraloperator/neuraloperator/pull/393), [#407](https://github.com/neuraloperator/neuraloperator/pull/407))
- The `DefaultDataProcessor` was updated to fix normalization issues in `training` and `eval` modes. ([397](https://github.com/neuraloperator/neuraloperator/pull/397))
- The Lp loss now has more sensible arguments names and defaults ([486](https://github.com/neuraloperator/neuraloperator/pull/486))
Contributors
This release was the result of a large team effort, thank you to all those who contributed!
dhpitt JeanKossaifi zongyi-li kovachki bonevbs mliuschi Robertboy18 animakumar
We particularly thank dhpitt, who joined the team as part of the Schmidt Academy for Software Engineering and has done a huge amount of work on this release.
And welcome to our new contributors, helgehr, suyashbire1, zijieli-Jlee, plutonium-239, ehsanngh and white-alistair!
**Full Changelog**: https://github.com/neuraloperator/neuraloperator/compare/0.3.0...1.0.0