| from pathlib import Path | |
| import sys,os,torch | |
| import torch.distributed as dist | |
| R=Path(__file__).resolve().parents[1];sys.path.insert(0,str(R));from model.pysteps 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') | |
| m=STEPS(**c['model']);p=R/c['paths']['checkpoint'] | |
| if rank==0:p.parent.mkdir(parents=True,exist_ok=True);torch.save({'model':m.state_dict(),'model_config':c['model']},p);write(R/c['paths']['training_metrics'],{'online_parameter_estimation':True,'world_size':world});print(p) | |
| if ddp:dist.destroy_process_group() | |