GRACE-SEDA / scripts /inference.py
zhangrenchao's picture
Publish GRACE-SEDA reproduction
bd6dbd0 verified
Raw
History Blame Contribute Delete
636 Bytes
from pathlib import Path
import sys,numpy as np,torch
R=Path(__file__).resolve().parents[1];sys.path.insert(0,str(R));from model.grace_seda import *
c=cfg(R);d=np.load(R/c['data']['path']);z=torch.load(R/c['paths']['checkpoint'],map_location='cpu',weights_only=True);x=torch.tensor(d['input']);pred=[]
with torch.no_grad():
for s in z['states']:m=GRACESEDA(**z['model_config']);m.load_state_dict(s);pred.append(m(x).numpy())
p=R/c['paths']['predictions'];p.parent.mkdir(parents=True,exist_ok=True);np.savez_compressed(p,ensemble=pred,grace=d['input'][:,0],wghm=d['input'][:,1],mean=np.mean(pred,0),uncertainty=np.std(pred,0));print(p)