Method test_load_weights

tests/bundle/test_bundle_download.py:313–361

(self, bundle_files, bundle_name, repo, device, model_file)

    @parameterized.expand([TEST_CASE_7])
    @skip_if_quick
    def test_load_weights(self, bundle_files, bundle_name, repo, device, model_file):
        with skip_if_downloading_fails():
            with tempfile.TemporaryDirectory() as tempdir:
                bundle_root = os.path.join(tempdir, bundle_name)
                # load weights
                model_1 = load(
                    name=bundle_name,
                    model_file=model_file,
                    bundle_dir=tempdir,
                    repo=repo,
                    source="github",
                    progress=False,
                    device=device,
                )
                # prepare network
                with open(os.path.join(bundle_root, bundle_files[2])) as f:
                    net_args = json.load(f)["network_def"]
                model_name = net_args["_target_"]
                del net_args["_target_"]
                model = getattr(nets, model_name)(**net_args)
                model.to(device)
                model.load_state_dict(model_1)
                model.eval()

                # prepare data and test
                input_tensor = torch.load(
                    os.path.join(bundle_root, bundle_files[4]), map_location=device, weights_only=True
                )
                output = model.forward(input_tensor)
                expected_output = torch.load(
                    os.path.join(bundle_root, bundle_files[3]), map_location=device, weights_only=True
                )
                assert_allclose(output, expected_output, atol=1e-3, rtol=1e-3, type_test=False)

                # load instantiated model directly and test, since the bundle has been downloaded,
                # there is no need to input `repo`
                _model_2 = getattr(nets, model_name)(**net_args)
                model_2 = load(
                    name=bundle_name,
                    model=_model_2,
                    model_file=model_file,
                    bundle_dir=tempdir,
                    progress=False,
                    device=device,
                    source="github",
                )
                model_2.eval()
                output_2 = model_2.forward(input_tensor)
                assert_allclose(output_2, expected_output, atol=1e-3, rtol=1e-3, type_test=False)

    @parameterized.expand([TEST_CASE_8])
    @skip_if_quick