Integrate Fourier Neural Operator (FNO) work

icenet-ai / icenet-roadmap

Central repository for the Turing-British Antarctic Survey IceNet Project

MIT License

7 stars 0 forks source link

Integrate Fourier Neural Operator (FNO) work #57

Open bnubald opened 8 months ago

bnubald commented 8 months ago

Build integration of @pyatsysh FNO work into the IceNet pipeline.

The goals are as follows:

[x] Get the current demo notebook running, figure out dependencies such that it runs.
[x] Get it running on BAS HPC, and create a reproducible environment (e.g. conda) that can be shared.
[x] Demonstrator notebook of JAX + UNet with the IceNet library/pipeline.
- https://github.com/icenet-ai/icenet-notebooks/tree/main/jax
[ ] Demonstrator notebook integration with the IceNet library/pipeline - using the IceNet dataloaders and output processing.
[ ] Integration with the icenet-pipeline for end to end runs.

@JimCircadian

bnubald commented 8 months ago

First task complete.

The second one is also mostly complete. I have run through the notebook on BAS HPC.

Still need to run through some updated dependency API changes with Orbax for checkpointing before completing the second task.

Will do so later this week.

bnubald commented 8 months ago

Added BAS HPC conda env file + nb run, and pushed commit to the working repo.

First two tasks are now complete.

Moving onto third task: Creation of a demonstrator notebook using IceNet dataloader and post processing (Need to use weekly averaging instead of daily inputs).

For initial implementation, will probably add it simply as an additional step after icenet data processing - for now). Although, do at some point need to allow monthly averaging to recreate the original paper.

bnubald commented 7 months ago

As a start to integrating JAX model development usage with IceNet, a demonstrator of Jax derived (Flax) UNet implementation (daily forecasting) using IceNet data download and plotting can be found here (Using Pytorch dataloader - deriving from work for #54 ).

The notebook also covers the usage of common JAX libraries such as Flax, Optax, Orbax, CommonLoopUtils to develop the model, train/validation/predict loops, usage of best checkpointing (based on validation rmse in this example) to save training states and load best one, custom weighted metrics, and more for model development.

Next, will be focusing on interface for FNO work, including pre-processing stages that are specific to this project.

bnubald commented 2 months ago

Links in with updates made (including torch dataloading) for #54, and pipeline integration architecture design with icenet v0.3.0_dev.