def test_valid_no_update(tmpdir):
p = tmpdir.join('model.pth.tar').strpath
net = FactorizedPrior(32, 64)
torch.save(net.state_dict(), p)
stdout, stderr = run_update_model(p, '--architecture', 'factorized-prior',
'--dir', tmpdir, '--no-update')
assert len(stdout) == 0
assert len(stderr) == 0
files = list(Path(tmpdir).glob('*.pth.tar'))
assert len(files) == 1
cdf_len = net.state_dict()['entropy_bottleneck._cdf_length']
new_cdf_len = torch.load(files[0])['entropy_bottleneck._cdf_length']
assert cdf_len.size(0) == new_cdf_len.size(0)
def test_invalid_model(tmpdir):
p = tmpdir.join("model.pth.tar").strpath
net = FactorizedPrior(32, 64)
torch.save(net.state_dict(), p)
with pytest.raises(SystemExit):
run_update_model(p, "--architecture", "foobar")
def test_valid(tmpdir):
p = tmpdir.join("model.pth.tar").strpath
net = FactorizedPrior(32, 64)
torch.save(net.state_dict(), p)
stdout, stderr = run_update_model(p, "--architecture", "factorized-prior",
"--dir", tmpdir)
assert len(stdout) == 0
assert len(stderr) == 0
files = list(Path(tmpdir).glob("*.pth.tar"))
assert len(files) == 1
cdf_len = net.state_dict()["entropy_bottleneck._cdf_length"]
new_cdf_len = torch.load(files[0])["entropy_bottleneck._cdf_length"]
assert cdf_len.size(0) != new_cdf_len.size(0)
def test_load(tmpdir):
p = tmpdir.join('model.pth.tar').strpath
net = FactorizedPrior(32, 64)
for k in ['network', 'state_dict']:
torch.save({k: net.state_dict()}, p)
stdout, stderr = run_update_model(p, '--architecture',
'factorized-prior', '--dir', tmpdir)
assert len(stdout) == 0
assert len(stderr) == 0
def test_load(tmpdir):
p = tmpdir.join("model.pth.tar").strpath
net = FactorizedPrior(32, 64)
for k in ["network", "state_dict"]:
torch.save({k: net.state_dict()}, p)
stdout, stderr = run_update_model(p, "--architecture",
"factorized-prior", "--dir", tmpdir)
assert len(stdout) == 0
assert len(stderr) == 0
def test_valid_name(tmpdir):
p = tmpdir.join('model.pth.tar').strpath
net = FactorizedPrior(32, 64)
torch.save(net.state_dict(), p)
stdout, stderr = run_update_model(p, '--architecture', 'factorized-prior',
'--dir', tmpdir, '--name', 'yolo')
assert len(stdout) == 0
assert len(stderr) == 0
files = sorted(list(Path(tmpdir).glob('*.pth.tar')))
assert len(files) == 2
assert files[0].name == 'model.pth.tar'
assert files[1].name[:5] == 'yolo-'
def test_valid_name(tmpdir):
p = tmpdir.join("model.pth.tar").strpath
net = FactorizedPrior(32, 64)
torch.save(net.state_dict(), p)
stdout, stderr = run_update_model(p, "--architecture", "factorized-prior",
"--dir", tmpdir, "--name", "yolo")
assert len(stdout) == 0
assert len(stderr) == 0
files = sorted(Path(tmpdir).glob("*.pth.tar"))
assert len(files) == 2
assert files[0].name == "model.pth.tar"
assert files[1].name[:5] == "yolo-"