OneScience-Group/TropiCycloneNet

1from pathlib import Path2import sys,os,math,torch3import torch.distributed as dist4from torch.nn.parallel import DistributedDataParallel as DDP5ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT))6from model.tropicyclonenet import *7c=load_config(ROOT);rank=int(os.getenv("RANK",0));world=int(os.getenv("WORLD_SIZE",1));distributed=world>18if distributed:dist.init_process_group("gloo")9torch.manual_seed(c["seed"]);base=TropiCycloneNet(**c["model"]);m=DDP(base) if distributed else base;opt=torch.optim.Adam(m.parameters(),lr=c["train"]["learning_rate"]);idx=list(range(rank,c["data"]["samples"],world));losses=[]10for _ in range(c["train"]["epochs"]):11 for i in idx:o,g,e,y,_=synthetic_sample(i);pred,prob=m(o[None],g[None],e[None]);errors=((pred-y[None,None])**2).mean((2,3));best=errors.min(1).values.mean();diversity=-pred.std(1).mean();loss=best+.01*diversity+.01*(-torch.log(prob+1e-8).mean());opt.zero_grad();loss.backward();opt.step();losses.append(float(loss))12total=torch.tensor([sum(losses),len(losses)],dtype=torch.float64)13if distributed:dist.all_reduce(total)14p=ROOT/c["paths"]["checkpoint"]15if rank==0:p.parent.mkdir(parents=True,exist_ok=True);torch.save({"model":base.state_dict(),"model_config":c["model"]},p);write_json(ROOT/c["paths"]["training_metrics"],{"loss":float(total[0]/total[1]),"world_size":world});print(p)16if distributed:dist.destroy_process_group()17