Request code for inference phase and the entire checkpoints for refining

[roxie62]

Dear authors,

Thank you very much for open-sourcing the code—I greatly enjoyed reading your paper! To facilitate further understanding of it, I want to ask if you might consider:

(1) Releasing the evaluation code used to generate predictions reported in the paper, and

(2) Uploading the full model checkpoints (including training iterations and optimizer states) to facilitate retraining/resuming training on new datasets?

Thank you very much for your help!

Many thanks,
Roxie

[mikemccabe210]

Hi Roxie,

Unfortunately, this is fairly old code at this point so what's released is what is going to be released. In terms of evals, the one-step metrics are already included in the repository - those can just be run on the test set of PDEBench to get the reported result. Any longer run metrics or viz are computed using autoregressive rollouts of the form (for one example - this is from some viz code):

def rollout(model, dset, state_labels, ic_index, steps=20, device='cpu'):
    model.eval()
    model.to(device)
    with torch.no_grad():
        x, bcs, y = dset[ic_index]
        targets = [torch.as_tensor(y)]
        preds = []
        x = torch.as_tensor(x).to(device).unsqueeze(1) 
        bcs = torch.as_tensor(bcs).unsqueeze(0).to(device)
        for i in range(1, steps+1):
            pred = model(x,  state_labels.to(device), bcs.to(device))
            preds.append(pred.squeeze(0))
            if i < steps:
                y = torch.as_tensor(dset[ic_index+i][-1])
                targets.append(y)
                x = torch.cat([x[1:], pred.unsqueeze(0)], 0)
        return list(map(lambda x: x.cpu(), preds))

with this being an example of how to invoke such a function from a viz notebook:

params = YParams(CONFIG_PATH, config_name=CONFIG_GROUP)
params['augmentation'] = False

model = build_avit(params) 
checkpoint = torch.load(CHECKPOINT_PATH, map_location=torch.device('cpu'))
model.load_state_dict(checkpoint['model_state'])
    
_, valid_dataset, _ = get_data_loader(params, params.valid_data_paths,
                                                                                 False,
                                                                        split='test')

def get_fields(valid_dataset, subset):
    return valid_dataset.subset_dict[subset.get_name()]

k = 0
subset = valid_dataset.sub_dsets[k]
indices = torch.as_tensor(get_fields(valid_dataset, subset)).to(device).unsqueeze(0)

preds = rollout(model, valid_dataset.sub_dsets[k], indices, 
                               0, steps=100, device=device)

I hope that helps!

Best,
Mike