File size: 5,950 Bytes
38e555d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | """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
# Four consecutive acquisitions x 4X4 pixels: all ten bands masked together.
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()
# Identical freshly initialized reconstruction heads in both arms; source encoder identical.
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.
# Materialize one fixed shuffled permutation for batch-boundary recovery.
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()
|