File size: 1,543 Bytes
05f7527 | 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 | import json,tempfile,hashlib
from pathlib import Path
import torch
from semantic_model import SemanticColorizer,save_semantic,load_semantic
from objectives import color_loss
torch.set_num_threads(2);torch.manual_seed(609)
results=[]
for head in ['dense','palette']:
m=SemanticColorizer(head=head,pretrained=True).eval()
count=sum(p.numel() for p in m.parameters());assert count<4_000_000
for shape in [(2,1,64,80),(1,1,173,241),(1,1,256,256)]:
with torch.no_grad():y=m(torch.zeros(shape))
assert y.shape==(shape[0],2,shape[2],shape[3]) and torch.isfinite(y).all()
x=torch.rand(2,1,64,64)*2-1;t=torch.zeros(2,2,64,64);t[:,0]=12;t[:,1]=18
opt=torch.optim.AdamW([p for n,p in m.named_parameters() if not n.startswith('encoder.')],lr=.001)
with torch.no_grad():initial=float(color_loss(m(x),t)[0])
for _ in range(16):
opt.zero_grad();loss,_=color_loss(m(x),t);loss.backward()
assert all(torch.isfinite(p.grad).all() for p in m.parameters() if p.grad is not None)
opt.step()
with torch.no_grad():final=float(color_loss(m(x),t)[0]);ref=m(x)
assert final<initial,(head,initial,final)
with tempfile.TemporaryDirectory() as tmp:
save_semantic(m,tmp);restored=load_semantic(tmp)
with torch.no_grad():torch.testing.assert_close(restored(x),ref)
results.append({'head':head,'parameters':count,'initial_loss':initial,'final_loss':final,'roundtrip':'passed','shape_tests':'passed','finite_backward':'passed'})
print('SMOKE_RESULTS',json.dumps(results),flush=True)
Path('smoke_results.json').write_text(json.dumps(results,indent=2))
|