Modulus Checkpoints: FourCastNet V2 Small

Modulus Checkpoints: FourCastNet V2 Small

Logo for Modulus Checkpoints: FourCastNet V2 Small
Description
FourCastNet V2 (small) model checkpoint package. FourCastNet V2 is a deep learning model for weather prediction that uses Spherical Harmonics Neural Operators.
Publisher
NVIDIA
Latest Version
v0.2
Modified
September 8, 2023
Size
2.93 GB

Details

This NGC asset is a FourCastNet v2 SFNO model checkpoint package. Model checkpoint package refers to the set of artifacts needed to run inference using pre-trained model which includes the model checkpoint, set of sample inputs, inference script.

Architecture

Please refer to the reference paper to learn about the model architecture. In training this checkpoint, Spherical Fourier Neural Operator (SFNO) is applied to forecasting atmospheric dynamics, and demonstrates stable autoregressive rollouts for a year of simulated time (1,460 steps), while retaining physically plausible dynamics. The SFNO has important implications for machine learning-based simulation of climate dynamics that could eventually help accelerate our response to climate change.

Training

The model is trained on a 73-channel subset of the ERA5 reanalysis data on single levels and pressure levels that is pre-processed and stored into HDF5 files.

How to use?

A minimal inference script is provided in this model checkpoint package to get you started easily.

To run inference on this checkpoint, you can follow the below steps (all the files needed are included in the zip file):

  1. Launch Modulus docker container

    docker run --rm --shm-size=1g --ulimit memlock=-1 --ulimit stack=67108864 --runtime nvidia -v ${PWD}:/examples -it nvcr.io/nvidia/modulus/modulus:23.08
    
  2. Download and install earth2mip.

    git clone https://github.com/NVIDIA/earth2mip.git
    cd earth2mip && pip install .   
    
  3. Download this checkpoint zip file, and unzip it

    wget 'https://api.ngc.nvidia.com/v2/models/nvidia/modulus/modulus_fcnv2_sm/versions/v0.2/files/fcnv2_sm.zip'
    unzip fcnv2_sm.zip
    
  4. Run the simple inference

    cd fcnv2_sm/
    python simple_inference.py
    

You can also use the model checkpoint from this model checkpoint package to evaluate it with the earth2mip framework.

References

Spherical Fourier Neural Operators: Learning Stable Dynamics on the Sphere