Skip to content

Commit

Permalink
provide pretrained checkpoints
Browse files Browse the repository at this point in the history
  • Loading branch information
tung-nd committed Feb 15, 2023
1 parent 99d0b4d commit d652c1b
Showing 1 changed file with 6 additions and 2 deletions.
8 changes: 6 additions & 2 deletions docs/usage.md
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,10 @@ python src/climax/pretrain/train.py --config configs/pretrain_climax.yaml \
!!! tip
Make sure to update the paths of the data directories in the config files (or override them via the CLI).

### Pretrained checkpoints
We provide two pretrained checkpoints, one was pretrained on [5.625deg](https://climaxrelease.blob.core.windows.net/checkpoints/ClimaX-5.625deg.ckpt) data, and the other was pretrained on [1.40625deg](https://climaxrelease.blob.core.windows.net/checkpoints/ClimaX-1.40625deg.ckpt) data. Both checkpoints were pretrained using all 5 CMIP6 datasets.
```Usage:``` We can load the checkpoint by passing the checkpoint url to the training script. See below.

## Global Forecasting

### Data Preparation
Expand Down Expand Up @@ -80,7 +84,7 @@ python src/climax/global_forecast/train.py --config configs/global_forecast_clim
--data.root_dir=/mnt/data/5.625deg_npz \
--data.predict_range=72 --data.out_variables=['z_500','t_850','t2m'] \
--data.batch_size=16 \
--model.pretrained_path=/mnt/checkpoints/climax_5.625deg.ckpt \
--model.pretrained_path='https://climaxrelease.blob.core.windows.net/checkpoints/ClimaX-5.625deg.ckpt' \
--model.lr=5e-7 --model.beta_1="0.9" --model.beta_2="0.99" \
--model.weight_decay=1e-5
```
Expand All @@ -107,7 +111,7 @@ python src/climax/regional_forecast/train.py --config configs/regional_forecast_
--data.region="NorthAmerica"
--data.predict_range=72 --data.out_variables=['z_500','t_850','t2m'] \
--data.batch_size=16 \
--model.pretrained_path=/mnt/checkpoints/climax_5.625deg.ckpt \
--model.pretrained_path='https://climaxrelease.blob.core.windows.net/checkpoints/ClimaX-5.625deg.ckpt' \
--model.lr=5e-7 --model.beta_1="0.9" --model.beta_2="0.99" \
--model.weight_decay=1e-5
```
Expand Down

0 comments on commit d652c1b

Please sign in to comment.