File size: 1,449 Bytes
4719196 | 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 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)
|