Repository: microsoft/qlib
File: qlib/rl/order_execution/policy.py (Python)

Code before:
            max_batchsize=max_batch_size,
            deterministic_eval=deterministic_eval,
            observation_space=obs_space,
            action_space=action_space,
        )
        if weight_file is not None:
            load_weight(self, weight_file)


# utilities: these should be put in a separate (common) file. #


def auto_device(module: nn.Module) -> torch.device:
    for param in module.parameters():
        return param.device
    return torch.device("cpu")  # fallback to cpu


def load_weight(policy: nn.Module, path: Path) -> None:
    assert isinstance(policy, nn.Module), "Policy has to be an nn.Module to load weight."
    loaded_weight = torch.load(path, map_location="cpu")

    if "vessel" in loaded_weight:
        loaded_weight = loaded_weight["vessel"]["policy"]
    try:
        policy.load_state_dict(loaded_weight)
    except RuntimeError:
        # try again by loading the converted weight
        # https://github.com/thu-ml/tianshou/issues/468
        for k in list(loaded_weight):
            loaded_weight["_actor_critic." + k] = loaded_weight[k]
        policy.load_state_dict(loaded_weight)


def chain_dedup(*iterables: Iterable) -> Generator[Any, None, None]:
    seen = set()
    for iterable in iterables:
        for i in iterable:
            if i not in seen:
                seen.add(i)
                yield i


Review comment (suggestion):
I think this should be handled by whoever calls `load_weight`.
For example, when the outer class receives a weight, it should first unpack it, and send the corresponding part to individual component.

This `if` can work as a hack though. Please leave a TODO here in that case.

Code after:
            max_batchsize=max_batch_size,
            deterministic_eval=deterministic_eval,
            observation_space=obs_space,
            action_space=action_space,
        )
        if weight_file is not None:
            load_weight(self, weight_file)


# utilities: these should be put in a separate (common) file. #


def auto_device(module: nn.Module) -> torch.device:
    for param in module.parameters():
        return param.device
    return torch.device("cpu")  # fallback to cpu


def load_weight(policy: nn.Module, path: Path) -> None:
    assert isinstance(policy, nn.Module), "Policy has to be an nn.Module to load weight."
    loaded_weight = torch.load(path, map_location="cpu")

    # TODO: this should be handled by whoever calls load_weight.
    # TODO: For example, when the outer class receives a weight, it should first unpack it,
    # TODO: and send the corresponding part to individual component.
    if "vessel" in loaded_weight:
        loaded_weight = loaded_weight["vessel"]["policy"]
    try:
        policy.load_state_dict(loaded_weight)
    except RuntimeError:
        # try again by loading the converted weight
        # https://github.com/thu-ml/tianshou/issues/468
        for k in list(loaded_weight):
            loaded_weight["_actor_critic." + k] = loaded_weight[k]
        policy.load_state_dict(loaded_weight)


def chain_dedup(*iterables: Iterable) -> Generator[Any, None, None]:
    seen = set()
    for iterable in iterables:
        for i in iterable:
            if i not in seen:
                seen.add(i)
                yield i