6834 Add MetaProterties and update load API

monai/bundle/__init__.py

@@ -13,10 +13,11 @@
 
 from .config_item import ComponentLocator, ConfigComponent, ConfigExpression, ConfigItem, Instantiable
 from .config_parser import ConfigParser
-from .properties import InferProperties, TrainProperties
+from .properties import InferProperties, MetaProterties, TrainProperties
 from .reference_resolver import ReferenceResolver
 from .scripts import (
     ckpt_export,
+    create_workflow,
     download,
     get_all_bundles_list,
     get_bundle_info,

monai/bundle/properties.py

@@ -13,7 +13,7 @@
 to interact with the bundle workflow.
 Some properties are required and some are optional, optional properties mean: if some component of the
 bundle workflow refer to the property, the property must be defined, otherwise, the property can be None.
-Every item in this `TrainProperties` or `InferProperties` dictionary is a property,
+Every item in this `TrainProperties` or `InferProperties` or `MetaProterties` dictionary is a property,
 the key is the property name and the values include:
 1. description.
 2. whether it's a required property.
@@ -48,6 +48,11 @@
         BundleProperty.REQUIRED: True,
         BundlePropertyConfig.ID: f"train{ID_SEP_KEY}trainer",
     },
+    "network_def": {
+        BundleProperty.DESC: "network module for the training.",
+        BundleProperty.REQUIRED: False,
+        BundlePropertyConfig.ID: "network_def",
+    },
     "max_epochs": {
         BundleProperty.DESC: "max number of epochs to execute the training.",
         BundleProperty.REQUIRED: True,
@@ -216,3 +221,42 @@
         BundlePropertyConfig.REF_ID: f"evaluator{ID_SEP_KEY}key_val_metric",
     },
 }
+
+MetaProterties = {
+    "version": {
+        BundleProperty.DESC: "bundle version",
+        BundleProperty.REQUIRED: False,
+        BundlePropertyConfig.ID: f"_meta_{ID_SEP_KEY}version",
+    },
+    "monai_version": {
+        BundleProperty.DESC: "required monai version used for bundle",
+        BundleProperty.REQUIRED: False,
+        BundlePropertyConfig.ID: f"_meta_{ID_SEP_KEY}monai_version",
+    },
+    "pytorch_version": {
+        BundleProperty.DESC: "required pytorch version used for bundle",
+        BundleProperty.REQUIRED: False,
+        BundlePropertyConfig.ID: f"_meta_{ID_SEP_KEY}pytorch_version",
+    },
+    "numpy_version": {
+        BundleProperty.DESC: "required numpy version used for bundle",
+        BundleProperty.REQUIRED: False,
+        BundlePropertyConfig.ID: f"_meta_{ID_SEP_KEY}numpy_version",
+    },
+    "description": {
+        BundleProperty.DESC: "description for bundle",
+        BundleProperty.REQUIRED: False,
+        BundlePropertyConfig.ID: f"_meta_{ID_SEP_KEY}description",
+    },
+    "spatial_shape": {
+        BundleProperty.DESC: "spatial shape for the inputs",
+        BundleProperty.REQUIRED: False,
+        BundlePropertyConfig.ID: f"_meta_{ID_SEP_KEY}network_data_format{ID_SEP_KEY}inputs{ID_SEP_KEY}image"
+        f"{ID_SEP_KEY}spatial_shape",
+    },
+    "channel_def": {
+        BundleProperty.DESC: "channel definition for the prediction",
+        BundleProperty.REQUIRED: False,
+        BundlePropertyConfig.ID: f"_meta_{ID_SEP_KEY}network_data_format{ID_SEP_KEY}outputs{ID_SEP_KEY}pred{ID_SEP_KEY}channel_def",
+    },
+}

monai/bundle/scripts.py

@@ -28,7 +28,6 @@
 
 from monai.apps.mmars.mmars import _get_all_ngc_models
 from monai.apps.utils import _basename, download_url, extractall, get_logger
-from monai.bundle.config_item import ConfigComponent
 from monai.bundle.config_parser import ConfigParser
 from monai.bundle.utils import DEFAULT_INFERENCE, DEFAULT_METADATA
 from monai.bundle.workflows import BundleWorkflow, ConfigWorkflow
@@ -63,7 +62,7 @@
 
 # set BUNDLE_DOWNLOAD_SRC="ngc" to use NGC source in default for bundle download
 # set BUNDLE_DOWNLOAD_SRC="monaihosting" to use monaihosting source in default for bundle download
-download_source = os.environ.get("BUNDLE_DOWNLOAD_SRC", "github")
+DEFAULT_DOWNLOAD_SOURCE = os.environ.get("BUNDLE_DOWNLOAD_SRC", "github")
 PPRINT_CONFIG_N = 5
 
 
@@ -253,7 +252,7 @@
     name: str | None = None,
     version: str | None = None,
     bundle_dir: PathLike | None = None,
-    source: str = download_source,
+    source: str = DEFAULT_DOWNLOAD_SOURCE,
     repo: str | None = None,
     url: str | None = None,
     remove_prefix: str | None = "monai_",
@@ -378,19 +377,26 @@
 
 def load(
     name: str,
+    model: torch.nn.Module | None = None,
     version: str | None = None,
+    workflow: str = "train",
     model_file: str | None = None,
     load_ts_module: bool = False,
     bundle_dir: PathLike | None = None,
-    source: str = download_source,
+    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] = (),
-    net_name: str | None = None,
-    **net_kwargs: Any,
+    dst_prefix: str = "",
+    mapping: dict | None = None,
+    exclude_vars: str | None = None,
+    inplace: bool = True,
+    workflow_name: str | BundleWorkflow | None = None,
+    args_file: str | None = None,
+    **net_override: Any,
 ) -> object | tuple[torch.nn.Module, dict, dict] | Any:
     """
     Load model weights or TorchScript module of a bundle.
@@ -402,8 +408,15 @@
             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.
+        workflow: 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.
@@ -417,37 +430,44 @@
             If used, it should be in the form of "repo_owner/repo_name/release_tag".
         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 two sources, remove prefix is necessary.
+            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
             weights. if not nested checkpoint, no need to set.
         config_files: extra filenames would be loaded. The argument only works when loading a TorchScript module,
             see `_extra_files` in `torch.jit.load` for more details.
-        net_name: if not `None`, a corresponding network will be instantiated and load the achieved weights.
-            This argument only works when loading weights.
-        net_kwargs: other arguments that are used to instantiate the network class defined by `net_name`.
+        dst_prefix: used in `copy_model_state`, `dst` key prefix, so that `dst[dst_prefix + src_key]`
+            will be assigned to the value of `src[src_key]`.
+        mapping: used in `copy_model_state`, a `{"src_key": "dst_key"}` dict, indicating that `dst[dst_prefix + dst_key]`
+            to be assigned to the value of `src[src_key]`.
+        exclude_vars: used in `copy_model_state`, a regular expression to match the `dst` variable names,
+            so that their values are not overwritten by `src`.
+        inplace: used in `copy_model_state`, whether to set the `dst` module with the updated `state_dict` via `load_state_dict`.
+            This option is only available when `dst` is a `torch.nn.Module`.
+        workflow_name: specified bundle workflow name, should be a string or class, default to "ConfigWorkflow".
+        args_file: a JSON or YAML file to provide default values for all the args in "download" function.
+        net_override: id-value pairs to override the parameters in the network of the bundle.
 
     Returns:
-        1. If `load_ts_module` is `False` and `net_name` is `None`, return model weights.
-        2. If `load_ts_module` is `False` and `net_name` is not `None`,
+        1. If `load_ts_module` is `False` and `model` is `None`,
+            return model weights if can't find "network_def" in the bundle,
+            else return an instantiated network that loaded the weights.
+        2. If `load_ts_module` is `False` and `model` is not `None`,
             return an instantiated network that loaded the weights.
         3. If `load_ts_module` is `True`, return a triple that include a TorchScript module,
             the corresponding metadata dict, and extra files dict.
             please check `monai.data.load_net_with_metadata` for more details.
 
     """
     bundle_dir_ = _process_bundle_dir(bundle_dir)
-
+    if device is None:
+        device = "cuda:0" if is_available() else "cpu"
     if model_file is None:
         model_file = os.path.join("models", "model.ts" if load_ts_module is True else "model.pt")
-    if source == "ngc":
-        name = _add_ngc_prefix(name)
-        if remove_prefix:
-            name = _remove_ngc_prefix(name, prefix=remove_prefix)
     full_path = os.path.join(bundle_dir_, name, model_file)
-    if not os.path.exists(full_path):
+    if not os.path.exists(full_path) or model is None:
         download(
             name=name,
             version=version,
@@ -456,10 +476,21 @@
             repo=repo,
             remove_prefix=remove_prefix,
             progress=progress,
+            args_file=args_file,
         )
+        train_config_file = bundle_dir_ / name / "configs" / f"{workflow}.json"
+        if train_config_file.is_file():
+            _net_override = {f"network_def#{key}": value for key, value in net_override.items()}
+            _workflow = create_workflow(
+                workflow_name=workflow_name,
+                args_file=args_file,
+                config_file=str(train_config_file),
+                workflow=workflow,
+                **_net_override,
+            )
+        else:
+            _workflow = None
 
-    if device is None:
-        device = "cuda:0" if is_available() else "cpu"
     # loading with `torch.jit.load`
     if load_ts_module is True:
         return load_net_with_metadata(full_path, map_location=torch.device(device), more_extra_files=config_files)
@@ -469,13 +500,19 @@
         warnings.warn(f"the state dictionary from {full_path} should be a dictionary but got {type(model_dict)}.")
         model_dict = get_state_dict(model_dict)
 
-    if net_name is None:
+    if model is None and _workflow is None:
         return model_dict
-    net_kwargs["_target_"] = net_name
-    configer = ConfigComponent(config=net_kwargs)
-    model = configer.instantiate()
-    model.to(device)  # type: ignore
-    copy_model_state(dst=model, src=model_dict if key_in_ckpt is None else model_dict[key_in_ckpt])  # type: ignore
+    model = _workflow.network_def if model is None else model  # type: ignore
+    model.to(device)
+
+    copy_model_state(
+        dst=model,
+        src=model_dict if key_in_ckpt is None else model_dict[key_in_ckpt],
+        dst_prefix=dst_prefix,
+        mapping=mapping,
+        exclude_vars=exclude_vars,
+        inplace=inplace,
+    )
     return model
 
 
@@ -675,12 +712,12 @@
         final_id: ID name of the expected config expression to finalize after running, default to "finalize".
             it's optional for both configs and this `run` function.
         meta_file: filepath of the metadata file, if it is a list of file paths, the content of them will be merged.
-            Default to "configs/metadata.json", which is commonly used for bundles in MONAI model zoo.
+            Default to None.
         config_file: filepath of the config file, if `None`, must be provided in `args_file`.
             if it is a list of file paths, the content of them will be merged.
         logging_file: config file for `logging` module in the program. for more details:
             https://docs.python.org/3/library/logging.config.html#logging.config.fileConfig.
-            Default to "configs/logging.conf", which is commonly used for bundles in MONAI model zoo.
+            Default to None.
         tracking: if not None, enable the experiment tracking at runtime with optionally configurable and extensible.
             if "mlflow", will add `MLFlowHandler` to the parsed bundle with default tracking settings,
             if other string, treat it as file path to load the tracking settings.
@@ -712,11 +749,11 @@
     config_file_, meta_file_, init_id_, run_id_, final_id_, logging_file_, tracking_ = _pop_args(
         _args,
         config_file=None,
-        meta_file="configs/metadata.json",
+        meta_file=None,
         init_id="initialize",
         run_id="run",
         final_id="finalize",
-        logging_file="configs/logging.conf",
+        logging_file=None,
         tracking=None,
     )
     workflow = ConfigWorkflow(
@@ -734,7 +771,9 @@
     workflow.finalize()
 
 
-def run_workflow(workflow: str | BundleWorkflow | None = None, args_file: str | None = None, **kwargs: Any) -> None:
+def run_workflow(
+    workflow_name: str | BundleWorkflow | None = None, args_file: str | None = None, **kwargs: Any
+) -> None:
     """
     Specify `bundle workflow` to run monai bundle components and workflows.
     The workflow should be subclass of `BundleWorkflow` and be available to import.
@@ -748,35 +787,17 @@
         python -m monai.bundle run_workflow --meta_file <meta path> --config_file <config path>
 
         # Set the workflow to other customized BundleWorkflow subclass:
-        python -m monai.bundle run_workflow --workflow CustomizedWorkflow ...
+        python -m monai.bundle run_workflow --workflow_name CustomizedWorkflow ...
 
     Args:
-        workflow: specified bundle workflow name, should be a string or class, default to "ConfigWorkflow".
+        workflow_name: specified bundle workflow name, should be a string or class, default to "ConfigWorkflow".
         args_file: a JSON or YAML file to provide default values for this API.
             so that the command line inputs can be simplified.
         kwargs: arguments to instantiate the workflow class.
 
     """
 
-    _args = _update_args(args=args_file, workflow=workflow, **kwargs)
-    _log_input_summary(tag="run", args=_args)
-    (workflow_name,) = _pop_args(_args, workflow=ConfigWorkflow)  # the default workflow name is "ConfigWorkflow"
-    if isinstance(workflow_name, str):
-        workflow_class, has_built_in = optional_import("monai.bundle", name=str(workflow_name))  # search built-in
-        if not has_built_in:
-            workflow_class = locate(str(workflow_name))  # search dotted path
-        if workflow_class is None:
-            raise ValueError(f"cannot locate specified workflow class: {workflow_name}.")
-    elif issubclass(workflow_name, BundleWorkflow):
-        workflow_class = workflow_name
-    else:
-        raise ValueError(
-            "Argument `workflow` must be a bundle workflow class name"
-            f"or subclass of BundleWorkflow, got: {workflow_name}."
-        )
-
-    workflow_ = workflow_class(**_args)
-    workflow_.initialize()
+    workflow_ = create_workflow(workflow_name=workflow_name, args_file=args_file, **kwargs)
     workflow_.run()
     workflow_.finalize()
 
@@ -1539,3 +1560,55 @@
         copyfile(str(ckpt_file), str(models_dir / "model.pt"))
     elif network is not None:
         save_state(network, str(models_dir / "model.pt"))
+
+
+def create_workflow(
+    workflow_name: str | BundleWorkflow | None = None,
+    config_file: str | Sequence[str] | None = None,
+    args_file: str | None = None,
+    **kwargs: Any,
+) -> Any:
+    """
+    Specify `bundle workflow` to create monai bundle workflows.
+    The workflow should be subclass of `BundleWorkflow` and be available to import.
+    It can be MONAI existing bundle workflows or user customized workflows.
+
+    Typical usage examples:
+
+    .. code-block:: python
+
+        # Specify config_file path to create workflow:
+        workflow = create_workflow(config_file="/workspace/spleen_ct_segmentation/configs/train.json", workflow="train")
+
+        # Set the workflow to other customized BundleWorkflow subclass to create workflow:
+        workflow = create_workflow(workflow_name=CustomizedWorkflow)
+
+    Args:
+        workflow_name: specified bundle workflow name, should be a string or class, default to "ConfigWorkflow".
+        config_file: filepath of the config file, if it is a list of file paths, the content of them will be merged.
+        args_file: a JSON or YAML file to provide default values for this API.
+            so that the command line inputs can be simplified.
+        kwargs: arguments to instantiate the workflow class.
+
+    """
+    _args = _update_args(args=args_file, workflow_name=workflow_name, config_file=config_file, **kwargs)
+    _log_input_summary(tag="run", args=_args)
+    (workflow_name,) = _pop_args(_args, workflow_name=ConfigWorkflow)  # the default workflow name is "ConfigWorkflow"
+    if isinstance(workflow_name, str):
+        workflow_class, has_built_in = optional_import("monai.bundle", name=str(workflow_name))  # search built-in
+        if not has_built_in:
+            workflow_class = locate(str(workflow_name))  # search dotted path
+        if workflow_class is None:
+            raise ValueError(f"cannot locate specified workflow class: {workflow_name}.")
+    elif issubclass(workflow_name, BundleWorkflow):  # type: ignore
+        workflow_class = workflow_name
+    else:
+        raise ValueError(
+            "Argument `workflow` must be a bundle workflow class name"
+            f"or subclass of BundleWorkflow, got: {workflow_name}."
+        )
+
+    workflow_ = workflow_class(**_args)
+    workflow_.initialize()
+
+    return workflow_