Function _replace_modules

monai/networks/utils.py:1093–1128

Helper function for :py:class:`monai.networks.utils.replace_modules`.

(
    parent: torch.nn.Module,
    name: str,
    new_module: torch.nn.Module,
    out: list[tuple[str, torch.nn.Module]],
    strict_match: bool = True,
    match_device: bool = True,
)

def _replace_modules(
    parent: torch.nn.Module,
    name: str,
    new_module: torch.nn.Module,
    out: list[tuple[str, torch.nn.Module]],
    strict_match: bool = True,
    match_device: bool = True,
) -> None:
    """
    Helper function for :py:class:`monai.networks.utils.replace_modules`.
    """
    if match_device:
        devices = list({i.device for i in parent.parameters()})
        # if only one device for whole of model
        if len(devices) == 1:
            new_module.to(devices[0])
    idx = name.find(".")
    # if there is "." in name, call recursively
    if idx != -1:
        parent_name = name[:idx]
        parent = getattr(parent, parent_name)
        name = name[idx + 1 :]
        _out: list[tuple[str, torch.nn.Module]] = []
        _replace_modules(parent, name, new_module, _out)
        # prepend the parent name
        out += [(f"{parent_name}.{r[0]}", r[1]) for r in _out]
    # no "." in module name, do the actual replacing
    else:
        if strict_match:
            old_module = getattr(parent, name)
            setattr(parent, name, new_module)
            out += [(name, old_module)]
        else:
            for mod_name, _ in parent.named_modules():
                if name in mod_name:
                    _replace_modules(parent, mod_name, deepcopy(new_module), out, strict_match=True)


def replace_modules(