OneScience-Group/Scale-MAE

1"""Train Scale-MAE with DDP, AMP, accumulation, warmup and cosine decay."""2import argparse, importlib.util, json, os, random3from pathlib import Path4import numpy as np, torch, yaml5from torch import distributed as dist6from torch.nn.parallel import DistributedDataParallel7from torch.utils.data import DataLoader, Dataset, DistributedSampler8ROOT=Path(__file__).resolve().parents[1]9class NPZDataset(Dataset):10    def __init__(self,path):11        a=np.load(path); self.images,self.targets,self.gsd=a["images"],a["targets"],a["gsd"]12    def __len__(self): return len(self.images)13    def __getitem__(self,i): return torch.from_numpy(self.images[i]), torch.from_numpy(self.targets[i]), torch.tensor(self.gsd[i], dtype=torch.float32)14def model_class():15    spec=importlib.util.spec_from_file_location("scalemae",ROOT/"model/scalemae.py"); mod=importlib.util.module_from_spec(spec); spec.loader.exec_module(mod); return mod.ScaleMAE16def main():17    p=argparse.ArgumentParser(); p.add_argument("--config",type=Path,default=ROOT/"conf/config.yaml"); p.add_argument("--data",type=Path); p.add_argument("--output",type=Path); p.add_argument("--epochs",type=int); p.add_argument("--batch-size",type=int); p.add_argument("--device",choices=("auto","cpu","cuda")); a=p.parse_args(); cfg=yaml.safe_load(a.config.read_text())18    if a.epochs is not None: cfg["training"]["epochs"] = a.epochs19    if a.batch_size is not None: cfg["training"]["batch_size"] = a.batch_size20    world=int(os.environ.get("WORLD_SIZE","1")); rank=int(os.environ.get("RANK","0")); local=int(os.environ.get("LOCAL_RANK","0")); requested=a.device or cfg["runtime"]["device"]21    if requested == "cuda" and not torch.cuda.is_available(): raise RuntimeError("CUDA requested but unavailable")22    cuda=torch.cuda.is_available() and requested!="cpu"; device=torch.device(f"cuda:{local}" if cuda else "cpu")23    if cuda: torch.cuda.set_device(local)24    if world>1: dist.init_process_group("nccl" if cuda else "gloo")25    seed=cfg["seed"]+rank; random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)26    source=a.data or ROOT/cfg["data"]["root"]/"train.npz"27    if not source.exists(): raise FileNotFoundError("Run scripts/fake_data.py first")28    ds=NPZDataset(source); sampler=DistributedSampler(ds,num_replicas=world,rank=rank) if world>1 else None; loader=DataLoader(ds,batch_size=cfg["training"]["batch_size"],shuffle=sampler is None,sampler=sampler,num_workers=cfg["training"]["num_workers"])29    model=model_class()(**cfg["model"]).to(device); model=DistributedDataParallel(model,device_ids=[local]) if world>1 else model; base=model.module if hasattr(model,"module") else model30    opt=torch.optim.AdamW(model.parameters(),lr=cfg["training"]["learning_rate"],betas=(.9,.95),weight_decay=cfg["training"]["weight_decay"]); steps=max(1,cfg["training"]["epochs"]*((len(ds)+cfg["training"]["batch_size"]-1)//cfg["training"]["batch_size"])); warm=max(1,int(steps*cfg["training"]["warmup_fraction"])); accum=cfg["training"]["gradient_accumulation"]31    amp=torch.amp.GradScaler("cuda", enabled=cuda and cfg["training"]["amp"]); history=[]; opt.zero_grad(set_to_none=True); step=032    for epoch in range(cfg["training"]["epochs"]):33        if sampler: sampler.set_epoch(epoch)34        model.train(); totals=torch.zeros(3,device=device)35        for batch_idx,(images,targets,gsd) in enumerate(loader):36            with torch.autocast(device_type="cuda",enabled=cuda and cfg["training"]["amp"]): out=model(images.to(device),gsd.to(device),target=targets.to(device)); loss=out["loss"]/accum37            if not torch.isfinite(loss): raise FloatingPointError("non-finite training loss")38            amp.scale(loss).backward(); totals += torch.stack([out["loss"].detach(),out["low_loss"].detach(),out["high_loss"].detach()])39            if (batch_idx+1)%accum==0 or batch_idx+1==len(loader):40                amp.unscale_(opt); torch.nn.utils.clip_grad_norm_(model.parameters(),1.0); amp.step(opt); amp.update(); opt.zero_grad(set_to_none=True); step+=1; lr=cfg["training"]["learning_rate"]*(step/warm if step<=warm else .5*(1+np.cos(np.pi*(step-warm)/max(1,steps-warm)))); [g.update(lr=lr) for g in opt.param_groups]41        if world>1: dist.all_reduce(totals); totals/=world42        totals/=max(len(loader),1); history.append({"epoch":epoch+1,"loss":totals[0].item(),"low_frequency_loss":totals[1].item(),"high_frequency_loss":totals[2].item()})43        if rank==0: print(history[-1])44    if rank==0:45        ck=a.output or ROOT/cfg["paths"]["checkpoint"]; met=ROOT/cfg["paths"]["training_metrics"]; ck.parent.mkdir(parents=True,exist_ok=True); met.parent.mkdir(parents=True,exist_ok=True); torch.save({"model":base.state_dict(),"config":cfg,"optimizer":opt.state_dict(),"epoch":cfg["training"]["epochs"],"seed":cfg["seed"]},ck); met.write_text(json.dumps({"history":history,"protocol":cfg["data"]["protocol"],"data_source":"synthetic"},indent=2)+"\n"); print("checkpoint=",ck)46    if world>1: dist.destroy_process_group()47if __name__=="__main__": main()48