| """Matched one-epoch fold5 continuation pilots; never overwrite supplied weights.""" |
| import argparse,time |
| from pathlib import Path |
| import torch |
| import torch.nn.functional as F |
| from torch.utils.data import DataLoader |
| import config as C |
| from dataset import PASTISPatchDataset,pretrain_patch_ids,pad_collate,crop_ij |
| from model import UTAE_ASPP,UTAEPrediction |
| from temporal_cnn import valid_from_index_positions |
| from notebook_helpers import write_json |
|
|
| class StructuredPrediction(UTAEPrediction): |
| def forward(self,x,pos): |
| b,t,c,h,w=x.shape |
| |
| coarse=torch.rand(b,(t+3)//4,(h+3)//4,(w+3)//4,device=x.device)<self.mask_ratio |
| hidden=coarse.repeat_interleave(4,1).repeat_interleave(4,2).repeat_interleave(4,3)[:,:t,:h,:w] |
| visible=(~hidden)[:,:,None].expand_as(x) |
| masked=x*visible+self.mask_token*(~visible) |
| out,spatial=self.utae(masked,pos) |
| recon=self.linear(out).permute(0,3,4,1,2) |
| recon2=self.midlinear(spatial.permute(0,2,3,1)).permute(0,3,1,2).reshape(b,t,c,h,w) |
| return recon,recon2,x,visible.float() |
|
|
| def main(): |
| ap=argparse.ArgumentParser();ap.add_argument('--pretrain_pth',required=True);ap.add_argument('--epochs',type=int,default=3);ap.add_argument('--policy',choices=['inherited','structured'],required=True);ap.add_argument('--limit_batches',type=int,default=0);ap.add_argument('--out_dir',required=True);args=ap.parse_args() |
| torch.set_num_threads(4);torch.manual_seed(3407);torch.backends.cuda.matmul.allow_tf32=True;torch.backends.cudnn.allow_tf32=True |
| out=Path(args.out_dir);out.mkdir(parents=True,exist_ok=True) |
| if (out/'completed.json').exists():return |
| enc=UTAE_ASPP(n_channels=C.N_CHANNELS,d_model=C.D_MODEL,n_heads=C.N_HEADS) |
| enc.load_state_dict(torch.load(args.pretrain_pth,map_location='cpu',weights_only=False)['model_state_dict'],strict=True) |
| model=(StructuredPrediction if args.policy=='structured' else UTAEPrediction)(enc,mask_ratio=.4).cuda() |
| |
| dl=DataLoader(PASTISPatchDataset(pretrain_patch_ids(),load_target=False),batch_size=4,shuffle=True,num_workers=2,pin_memory=True,persistent_workers=True,collate_fn=pad_collate) |
| opt=torch.optim.AdamW(model.parameters(),lr=1e-5,betas=(.9,.95),weight_decay=C.PRE_WD);scaler=torch.amp.GradScaler('cuda') |
| start_epoch=0;start_batch=0;steps=0;total=0.;hidden_count=valid_count=0. |
| |
| generator=torch.Generator().manual_seed(3407);order=torch.randperm(496,generator=generator).tolist() |
| dl=DataLoader(torch.utils.data.Subset(dl.dataset,order),batch_size=4,shuffle=False,num_workers=2,pin_memory=True,persistent_workers=True,collate_fn=pad_collate) |
| if (out/'resume.tar').exists(): |
| ck=torch.load(out/'resume.tar',weights_only=False);model.load_state_dict(ck['model']);opt.load_state_dict(ck['optimizer']);scaler.load_state_dict(ck['scaler']);start_epoch=ck.get('epoch',0);start_batch=ck['next_batch'];steps=ck['steps'];total=ck['total'];hidden_count=ck['hidden_count'];valid_count=ck['valid_count'];torch.set_rng_state(ck['rng']);torch.cuda.set_rng_state_all(ck['cuda_rng']) |
| started=time.monotonic();model.train() |
| for epoch in range(start_epoch,args.epochs): |
| if epoch != start_epoch: |
| start_batch=0;steps=0;total=0.;hidden_count=valid_count=0. |
| for bi,((x,pos,_),_) in enumerate(dl): |
| if bi<start_batch:continue |
| if args.limit_batches and bi>=args.limit_batches:break |
| x,pos=x.cuda(),pos.cuda();valid=valid_from_index_positions(pos)[:,:,None,None,None] |
| for i in range(4): |
| for j in range(4): |
| xc,_=crop_ij(x,None,i,j,4);opt.zero_grad(set_to_none=True) |
| with torch.autocast('cuda',dtype=torch.float16):r,r2,target,mask=model(xc,pos) |
| weight=(1-mask)*valid;denom=weight.sum().clamp_min(1.) |
| loss=(((r.float()-target)**2*weight).sum()+((r2.float()-target)**2*weight).sum())/denom |
| if not torch.isfinite(loss):raise RuntimeError('Nonfinite reconstruction loss') |
| scaler.scale(loss).backward();scaler.unscale_(opt);torch.nn.utils.clip_grad_norm_(model.parameters(),1.);scaler.step(opt);scaler.update() |
| steps+=1;total+=loss.item();hidden_count+=weight.sum().item();valid_count+=valid.expand_as(xc).sum().item() |
| ck={'model':model.state_dict(),'optimizer':opt.state_dict(),'scaler':scaler.state_dict(),'epoch':epoch,'next_batch':bi+1,'steps':steps,'total':total,'hidden_count':hidden_count,'valid_count':valid_count,'rng':torch.get_rng_state(),'cuda_rng':torch.cuda.get_rng_state_all()} |
| tmp=out/'resume.tmp';torch.save(ck,tmp);tmp.replace(out/'resume.tar') |
| write_json(out/'progress.json',{'epoch':epoch+1,'epochs':args.epochs,'batches_completed':bi+1,'batches':len(dl),'steps':steps,'loss':total/steps,'mask_fraction':hidden_count/max(valid_count,1),'seconds':time.monotonic()-started}) |
| if bi%5==0:print(args.policy,'epoch',epoch+1,'batch',bi+1,'/',len(dl),'steps',steps,'loss',total/steps,'mask_fraction',hidden_count/max(valid_count,1),'seconds',time.monotonic()-started,flush=True) |
| if not args.limit_batches:assert steps==1984 |
| torch.save({'epoch':args.epochs-1,'model_state_dict':model.utae.state_dict(),'policy':args.policy,'source':'latest.utae.tar','pilot_epochs':args.epochs},out/'encoder_32X32.tar') |
| write_json(out/'completed.json',{'policy':args.policy,'epochs':args.epochs,'total_steps':args.epochs*1984,'steps_last_epoch':steps,'mean_reconstruction_loss':total/steps,'mask_fraction':hidden_count/max(valid_count,1),'seconds':time.monotonic()-started,'pilot':True,'labels_used':False}) |
| if __name__=='__main__':main() |
|
|
