allenai/ACE2-ERA5-training-artifacts

ACE2-ERA5 training artifacts

This repository contains training artifacts generated when training the ACE2-ERA5 model for the ACE2 paper.

Specifically, it includes checkpoints from all four random seed training runs and some additional intermediate checkpoints from training.

These checkpoints may be useful for certain research purposes, but in general we recommend using the ACE2-ERA5 checkpoint from this repository, which was the model described in the ACE2 paper and which has the lowest climate bias of all these checkpoints.

The training artifacts are in directories named rsN where N ranges from 0 to 3. This corresponds to the four random seeds trained. Within each directory, the checkpoints available are:

• best_ckpt.tar (this is the checkpoint with lowest validation loss throughout training)
• best_inference_ckpt.tar (this is the checkpoint with lowest time-mean inference error through training, see Equation 8 of ACE2 paper)
• for rs2 and rs3, ckpt_0025.tar, ckpt_0050.tar, ckpt_0075.tar and ckpt_0100.tar are the checkpoints after 25, 50, 75 and 100 epochs of training

It is possible to check the number of training iterations that were performed for a particular checkpoint as follows:

python -c "import torch; print(torch.load('CKPT_FILENAME.tar')['num_batches_seen'])"

Note: the published ACE2-ERA5 checkpoint is rs2/best_inference_ckpt.tar.