File size: 1,196 Bytes
ae605ea | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 | from pathlib import Path
import sys,os,numpy as np,torch;import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
R=Path(__file__).resolve().parents[1];sys.path.insert(0,str(R));from model.fuxi_weather import *
c=cfg(R);rank=int(os.getenv('RANK',0));world=int(os.getenv('WORLD_SIZE',1));ddp=world>1
if ddp:dist.init_process_group('gloo')
d=np.load(R/c['data']['path']);base=FuXiWeather(**c['model']);m=DDP(base,find_unused_parameters=True) if ddp else base;opt=torch.optim.Adam(m.parameters(),lr=c['train']['learning_rate']);ls=[]
for i in range(rank,8,world):b,o,q,t=map(torch.tensor,(d['background'][i:i+1],d['observation'][i:i+1],d['mask'][i:i+1],d['target'][i:i+1]));a,f=m(b,o,q,i%2==1);loss=((a-t)**2).mean()+((f-t)**2).mean();opt.zero_grad();loss.backward();opt.step();ls.append(float(loss))
v=torch.tensor([sum(ls),len(ls)],dtype=torch.float64)
if ddp:dist.all_reduce(v)
p=R/c['paths']['checkpoint']
if rank==0:p.parent.mkdir(parents=True,exist_ok=True);torch.save({'model':base.state_dict(),'model_config':c['model']},p);write(R/c['paths']['training_metrics'],{'loss':float(v[0]/v[1]),'world_size':world});print(p)
if ddp:dist.destroy_process_group()
|