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