Function load

monai/bundle/scripts.py:628–768

Load model weights or TorchScript module of a bundle. Args: name: bundle name. If `None` and `url` is `None`, it must be provided in `args_file`. for example: "spleen_ct_segmentation", "prostate_mri_anatomy" in model-zoo: https://github.com/Proje

(
    name: str,
    model: torch.nn.Module | None = None,
    version: str | None = None,
    workflow_type: str = "train",
    model_file: str | None = None,
    load_ts_module: bool = False,
    bundle_dir: PathLike | None = None,
    source: str = DEFAULT_DOWNLOAD_SOURCE,
    repo: str | None = None,
    remove_prefix: str | None = "monai_",
    progress: bool = True,
    device: str | None = None,
    key_in_ckpt: str | None = None,
    config_files: Sequence[str] = (),
    workflow_name: str | BundleWorkflow | None = None,
    args_file: str | None = None,
    copy_model_args: dict | None = None,
    net_override: dict | None = None,
)

def load(
    name: str,
    model: torch.nn.Module | None = None,
    version: str | None = None,
    workflow_type: str = "train",
    model_file: str | None = None,
    load_ts_module: bool = False,
    bundle_dir: PathLike | None = None,
    source: str = DEFAULT_DOWNLOAD_SOURCE,
    repo: str | None = None,
    remove_prefix: str | None = "monai_",
    progress: bool = True,
    device: str | None = None,
    key_in_ckpt: str | None = None,
    config_files: Sequence[str] = (),
    workflow_name: str | BundleWorkflow | None = None,
    args_file: str | None = None,
    copy_model_args: dict | None = None,
    net_override: dict | None = None,
) -> object | tuple[torch.nn.Module, dict, dict] | Any:
    """
    Load model weights or TorchScript module of a bundle.

    Args:
        name: bundle name. If `None` and `url` is `None`, it must be provided in `args_file`.
            for example:
            "spleen_ct_segmentation", "prostate_mri_anatomy" in model-zoo:
            https://github.com/Project-MONAI/model-zoo/releases/tag/hosting_storage_v1.
            "monai_brats_mri_segmentation" in ngc:
            https://catalog.ngc.nvidia.com/models?filters=&orderBy=scoreDESC&query=monai.
            "mednist_gan" in monaihosting:
            https://api.ngc.nvidia.com/v2/models/nvidia/monaihosting/mednist_gan/versions/0.2.0/files/mednist_gan_v0.2.0.zip
        model: a pytorch module to be updated. Default to None, using the "network_def" in the bundle.
        version: version name of the target bundle to download, like: "0.1.0". If `None`, will download
            the latest version. If `source` is "huggingface_hub", this argument is a Git revision id.
        workflow_type: specifies the workflow type: "train" or "training" for a training workflow,
            or "infer", "inference", "eval", "evaluation" for a inference workflow,
            other unsupported string will raise a ValueError.
            default to `train` for training workflow.
        model_file: the relative path of the model weights or TorchScript module within bundle.
            If `None`, "models/model.pt" or "models/model.ts" will be used.
        load_ts_module: a flag to specify if loading the TorchScript module.
        bundle_dir: directory the weights/TorchScript module will be loaded from.
            Default is `bundle` subfolder under `torch.hub.get_dir()`.
        source: storage location name. This argument is used when `model_file` is not existing locally and need to be
            downloaded first.
            In default, the value is achieved from the environment variable BUNDLE_DOWNLOAD_SRC, and
            it should be "ngc", "monaihosting", "github", or "huggingface_hub".
        repo: repo name. This argument is used when `url` is `None` and `source` is "github" or "huggingface_hub".
            If `source` is "github", it should be in the form of "repo_owner/repo_name/release_tag".
            If `source` is "huggingface_hub", it should be in the form of "repo_owner/repo_name".
        remove_prefix: This argument is used when `source` is "ngc". Currently, all ngc bundles
            have the ``monai_`` prefix, which is not existing in their model zoo contrasts. In order to
            maintain the consistency between these three sources, remove prefix is necessary.
            Therefore, if specified, downloaded folder name will remove the prefix.
        progress: whether to display a progress bar when downloading.
        device: target device of returned weights or module, if `None`, prefer to "cuda" if existing.
        key_in_ckpt: for nested checkpoint like `{"model": XXX, "optimizer": XXX, ...}`, specify the key of model