This page highlights the steps to run the FourCastNet notebook in ARE and on local machine.

Code and Data Location

The datasets and pertained models for FourCastNet are available at /g/data/wb00/FourCastNet. You need to connect to dataset area wb00 to access it. A description of the dataset area (/g/data/wb00/FourCastNet) is provided below:

The example Jupyter Notebook may be found on the Github repository: https://github.com/nci/FourCastNet. The repository also contains a PBS script named 'job.sh', which may be used to run the inference scripts to generate the predictions.

Run the notebook in ARE

In order to run the notebook on ARE, you need to clone the FourCastNet repository to your own working space on Gadi. For example: If you are working on /scratch/ab12, then follow the given commands to clone FourCastNet.

cd /scratch/ab12/user_id
git clone https://github.com/nci/FourCastNet.git

Launch an ARE Instance

  1. Go to https://are.nci.org.au/ and login. 
  2. Click on the Jupyter notebook app and a form will appear.
  3. In the Dashboard for JupyterLab launch, fill up the followings.
    Walltime (hours): 1
    Queue: gpuvolta
    Compute Size: 1gpu

    Project: your own project
    Storage: gdata/dk92+gdata/wb00+gdata/rt52+(your working space storage)

  4. In the Advanced options section: 
    Module Directories: /g/data/dk92/apps/Modules/modulefiles
    Modules: NCI-ai-ml/23.03


  5. Now, Launch the JupyterLab Instance. 

Launch the Notebook

  1. In the jupyter notebook, open the notebook by navigating to /FourCastNet/notebooks/ and opening inference.ipynb.
  2. In the setting up section, remember to change the current directory by setting path in DIR = "/CURRENT/WORKING/DIRECTORY/".
     
  3. All the other code is already present and there is no need to write anything. If you want to run the cells, remember to start from top.

NOTE: If you wish against running the inference scripts in the notebook, you may directly use the output locations at /g/data/wb00/FourCastNet. The locations have been commented out in the notebook for your convenience. 

Run the PBS script on CLI

In order to run the script on the command line interface, you will need to first log in to Gadi using the following command:

ssh aaa777@gadi.nci.org.au

Replace aaa777 with your NCI username. 

In the FourCastNet directory, you can find a PBS script named job.sh. The script is designed to run the inference commands and produce the predictions similar to the notebook. 

#!/bin/bash

#PBS -q gpuvolta
#PBS -l ncpus=12
#PBS -l ngpus=1
#PBS -l jobfs=200GB
#PBS -l storage=gdata/dk92+gdata/wb00+(INSERT ANY ADDITIONAL STORAGE HERE)
#PBS -l mem=48GB
#PBS -l walltime=08:00:00
#PBS -P INSERT PROJECT

set -eu
  
module use /g/data/dk92/apps/Modules/modulefiles
module load NCI-ai-ml/23.03
module load cuda/11.7.0

pretrained_root=<pre-trained root>
output_path=<output path>

# Run inference AFNO Backbone 
python inference/inference.py \
    --config=afno_backbone \
    --run_num=0 --vis   \
    --weights '/g/data/wb00/FourCastNet/nvlab/v0/pretrained/backbone.ckpt' \
    --override_dir './output'

# Run inference_ensemble AFNO Backbone
python inference/inference_ensemble.py \
    --config=afno_backbone \
    --run_num=0 \
    --n_pert 10 \
    --override_dir './output'  \
    --weights '/g/data/wb00/FourCastNet/nvlab/v0/pretrained/backbone.ckpt'
    
# Run inference for precipitation
python inference/inference_precip.py \
    --config=precip \
    --run_num=0 \
    --vis \
    --weights '/g/data/wb00/FourCastNet/nvlab/v0/pretrained/precip.ckpt' \
    --override_dir './output'

# Run inference to generate preds.zarr
python -u inference_nci/inference.py \
  --start_time=2017-12-31T18:00:00 \
  --end_time=2018-12-30T18:00:00 \
  --checkpoint_dir=$pretrained_root/FCN_weights_v0 \
  --stats_dir=$pretrained_root/stats_v0 \
  --output_path=$output_path \
  --prediction_length=20

In the script, you will need to edit the #PBS -P directive to list your own project. You will also need to list any additional storage options if required in the #PBS -l storage directive. The output path should be set to the location where you want the output predictions to be stored. 






  • No labels