| import copy,json,sys,tempfile,unittest |
| from pathlib import Path |
| import numpy as np |
| import torch |
| ROOT=Path(__file__).resolve().parents[1] |
| sys.path.insert(0,str(ROOT/'code')) |
| from project_config import load_config |
| from train_recon import Net,load_dat |
| class ConfigTests(unittest.TestCase): |
| def test_checkpoint_compatible_structure(self): |
| model=Net();ch=[41,64,128,64,16,1];parts=[] |
| for i in range(5): |
| parts.append(torch.nn.Conv2d(ch[i],ch[i+1],3,1,1)) |
| if i<4:parts.append(torch.nn.ReLU(True)) |
| baseline=torch.nn.Sequential(*parts);baseline.load_state_dict(model.seq.state_dict()) |
| x=torch.randn(1,41,8,8) |
| torch.testing.assert_close(model(x),baseline(x),rtol=0,atol=0) |
| def test_configuration_changes_model_and_decoder(self): |
| c=load_config();c['input'].update(frames=3,height=4,width=8,packed_bytes=12,flip_height=False);c['frontend']['channels']=[3,5,1] |
| with tempfile.TemporaryDirectory() as d: |
| p=Path(d)/'config.json';p.write_text(json.dumps(c));raw=np.zeros((3,4,8),np.uint8);raw[0,0,0]=1;f=Path(d)/'x.dat';np.packbits(raw.reshape(3,-1),axis=1,bitorder='little').tofile(f) |
| np.testing.assert_array_equal(load_dat(f,p),raw) |
| self.assertEqual(Net(p)(torch.zeros(1,3,4,8)).shape,(1,1,4,8)) |
| def test_reject_invalid_config(self): |
| c=load_config();c['input']['packed_bytes']=1 |
| with tempfile.TemporaryDirectory() as d: |
| p=Path(d)/'config.json';p.write_text(json.dumps(c)) |
| with self.assertRaises(ValueError):load_config(p) |
|
|
