Method load_state_dict

cdslib/core/models/base_model.py:169–229

Load a pretrained model given a state dictionary.

(
        self,
        state_dict: T.Dict[str, T.Any],
    )

            v.eval()

    def load_state_dict(
        self,
        state_dict: T.Dict[str, T.Any],
    ):
        """Load a pretrained model given a state dictionary."""

        loaded_names = set()

        # gather all base models
        nets = inspect.getmembers(self, lambda x: isinstance(x, BaseModel))  # name, net
        for name, net in nets:
            if name in state_dict:
                getattr(self, name).load_state_dict(state_dict[name])
            else:
                warnings.warn(f"{name} not exist in checkpoint")
            loaded_names.add(name)

        # gather all nn.modules
        nets = inspect.getmembers(self, lambda x: isinstance(x, torch.nn.Module))  # name, net
        for name, net in nets:
            if name in state_dict:
                try:
                    getattr(self, name).load_state_dict(state_dict[name])
                except:
                    new_sub_state_dict = OrderedDict()
                    for key, val in state_dict[name].items():
                        if key.startswith("module."):
                            new_key = key[7:]
                            new_sub_state_dict[new_key] = val
                        else:
                            new_sub_state_dict[key] = val
                    getattr(self, name).load_state_dict(new_sub_state_dict)
            else:
                warnings.warn(f"{name} not exist in checkpoint")
            loaded_names.add(name)

        # load those in var_names_to_load other things
        buffer_names = state_dict.get("buffer_names", set())
        parameter_names = state_dict.get("parameter_names", set())

        var_names_to_save = state_dict.get("var_names_to_save", set())
        self.register_var_to_save(
            var_name=var_names_to_save,
        )

        var_names_to_load = state_dict.get("var_names_to_load", set())
        self.register_var_to_save(
            var_name=var_names_to_load,
        )

        # load only those in var_names_to_load
        for name in var_names_to_load - loaded_names:
            if name not in state_dict:
                continue
            setattr(self, name, state_dict[name])

            if name in buffer_names:
                self.register_buffer(name, state_dict[name], persistent=True)