| 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)) |
|
|
