| 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.grace_seda 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']);states=[];hist=[] |
| for member in range(5): |
| torch.manual_seed(c['seed']+member);base=GRACESEDA(**c['model']);m=DDP(base) if ddp else base;opt=torch.optim.Adam(m.parameters(),lr=c['train']['learning_rate']);ls=[] |
| for i in range(rank,20,world):x=torch.tensor(d['input'][i:i+1]);loss=loss_fn(m(x),x);opt.zero_grad();loss.backward();opt.step();ls.append(float(loss)) |
| if rank==0:states.append(base.state_dict());hist.append(float(np.mean(ls))) |
| p=R/c['paths']['checkpoint'] |
| if rank==0:p.parent.mkdir(parents=True,exist_ok=True);torch.save({'states':states,'model_config':c['model']},p);write(R/c['paths']['training_metrics'],{'ensemble_losses':hist,'world_size':world});print(p) |
| if ddp:dist.destroy_process_group() |
|
|
