def ckpt_export(
net_id: Optional[str] = None,
filepath: Optional[PathLike] = None,
ckpt_file: Optional[str] = None,
meta_file: Optional[Union[str, Sequence[str]]] = None,
config_file: Optional[Union[str, Sequence[str]]] = None,
key_in_ckpt: Optional[str] = None,
args_file: Optional[str] = None,
**override,
):
"""
Export the model checkpoint to the given filepath with metadata and config included as JSON files.
Typical usage examples:
.. code-block:: bash
python -m monai.bundle ckpt_export network --filepath <export path> --ckpt_file <checkpoint path> ...
Args:
net_id: ID name of the network component in the config, it must be `torch.nn.Module`.
filepath: filepath to export, if filename has no extension it becomes `.ts`.
ckpt_file: filepath of the model checkpoint to load.
meta_file: filepath of the metadata file, if it is a list of file paths, the content of them will be merged.
config_file: filepath of the config file to save in TorchScript model and extract network information,
the saved key in the TorchScript model is the config filename without extension, and the saved config
value is always serialized in JSON format no matter the original file format is JSON or YAML.
it can be a single file or a list of files. if `None`, must be provided in `args_file`.
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.
args_file: a JSON or YAML file to provide default values for `meta_file`, `config_file`,
`net_id` and override pairs. so that the command line inputs can be simplified.
override: id-value pairs to override or add the corresponding config content.
e.g. ``--_meta#network_data_format#inputs#image#num_channels 3``.
"""
_args = _update_args(
args=args_file,
net_id=net_id,
filepath=filepath,
meta_file=meta_file,
config_file=config_file,
ckpt_file=ckpt_file,
key_in_ckpt=key_in_ckpt,
**override,
)
_log_input_summary(tag="ckpt_export", args=_args)
filepath_, ckpt_file_, config_file_, net_id_, meta_file_, key_in_ckpt_ = _pop_args(
_args,
"filepath",
"ckpt_file",
"config_file",
net_id="",
meta_file=None,
key_in_ckpt="")
parser = ConfigParser()
parser.read_config(f=config_file_)
if meta_file_ is not None:
parser.read_meta(f=meta_file_)
# the rest key-values in the _args are to override config content
for k, v in _args.items():
parser[k] = v
net = parser.get_parsed_content(net_id_)
if has_ignite:
# here we use ignite Checkpoint to support nested weights and be compatible with MONAI CheckpointSaver
Checkpoint.load_objects(to_load={key_in_ckpt_: net},
checkpoint=ckpt_file_)
else:
copy_model_state(
dst=net,
src=ckpt_file_ if key_in_ckpt_ == "" else ckpt_file_[key_in_ckpt_])
# convert to TorchScript model and save with meta data, config content
net = convert_to_torchscript(model=net)
extra_files: Dict = {}
for i in ensure_tuple(config_file_):
# split the filename and directory
filename = os.path.basename(i)
# remove extension
filename, _ = os.path.splitext(filename)
if filename in extra_files:
raise ValueError(
f"filename '{filename}' is given multiple times in config file list."
)
extra_files[filename] = json.dumps(
ConfigParser.load_config_file(i)).encode()
save_net_with_metadata(
jit_obj=net,
filename_prefix_or_stream=filepath_,
include_config_vals=False,
append_timestamp=False,
meta_values=parser.get().pop("_meta_", None),
more_extra_files=extra_files,
)
logger.info(f"exported to TorchScript file: {filepath_}.")
def run(
runner_id: Optional[Union[str, Sequence[str]]] = None,
meta_file: Optional[Union[str, Sequence[str]]] = None,
config_file: Optional[Union[str, Sequence[str]]] = None,
logging_file: Optional[str] = None,
args_file: Optional[str] = None,
**override,
):
"""
Specify `meta_file` and `config_file` to run monai bundle components and workflows.
Typical usage examples:
.. code-block:: bash
# Execute this module as a CLI entry:
python -m monai.bundle run training --meta_file <meta path> --config_file <config path>
# Override config values at runtime by specifying the component id and its new value:
python -m monai.bundle run training --net#input_chns 1 ...
# Override config values with another config file `/path/to/another.json`:
python -m monai.bundle run evaluating --net %/path/to/another.json ...
# Override config values with part content of another config file:
python -m monai.bundle run training --net %/data/other.json#net_arg ...
# Set default args of `run` in a JSON / YAML file, help to record and simplify the command line.
# Other args still can override the default args at runtime:
python -m monai.bundle run --args_file "/workspace/data/args.json" --config_file <config path>
Args:
runner_id: ID name of the expected config expression to run, can also be a list of IDs to run in order.
meta_file: filepath of the metadata file, if it is a list of file paths, the content of them will be merged.
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, default to `None`. for more details:
https://docs.python.org/3/library/logging.config.html#logging.config.fileConfig.
args_file: a JSON or YAML file to provide default values for `runner_id`, `meta_file`,
`config_file`, `logging`, and override pairs. so that the command line inputs can be simplified.
override: id-value pairs to override or add the corresponding config content.
e.g. ``--net#input_chns 42``.
"""
_args = _update_args(
args=args_file,
runner_id=runner_id,
meta_file=meta_file,
config_file=config_file,
logging_file=logging_file,
**override,
)
if "config_file" not in _args:
raise ValueError(
f"`config_file` is required for 'monai.bundle run'.\n{run.__doc__}"
)
_log_input_summary(tag="run", args=_args)
config_file_, meta_file_, runner_id_, logging_file_ = _pop_args(
_args, "config_file", meta_file=None, runner_id="", logging_file=None)
if logging_file_ is not None:
logger.info(
f"set logging properties based on config: {logging_file_}.")
fileConfig(logging_file_, disable_existing_loggers=False)
parser = ConfigParser()
parser.read_config(f=config_file_)
if meta_file_ is not None:
parser.read_meta(f=meta_file_)
# the rest key-values in the _args are to override config content
for k, v in _args.items():
parser[k] = v
# resolve and execute the specified runner expressions in the config, return the results
return [
parser.get_parsed_content(i,
lazy=True,
eval_expr=True,
instantiate=True)
for i in ensure_tuple(runner_id_)
]
def verify_net_in_out(
net_id: Optional[str] = None,
meta_file: Optional[Union[str, Sequence[str]]] = None,
config_file: Optional[Union[str, Sequence[str]]] = None,
device: Optional[str] = None,
p: Optional[int] = None,
n: Optional[int] = None,
any: Optional[int] = None,
args_file: Optional[str] = None,
**override,
):
"""
Verify the input and output data shape and data type of network defined in the metadata.
Will test with fake Tensor data according to the required data shape in `metadata`.
Typical usage examples:
.. code-block:: bash
python -m monai.bundle verify_net_in_out network --meta_file <meta path> --config_file <config path>
Args:
net_id: ID name of the network component to verify, it must be `torch.nn.Module`.
meta_file: filepath of the metadata file to get network args, if `None`, must be provided in `args_file`.
if it is a list of file paths, the content of them will be merged.
config_file: filepath of the config file to get network definition, if `None`, must be provided in `args_file`.
if it is a list of file paths, the content of them will be merged.
device: target device to run the network forward computation, if None, prefer to "cuda" if existing.
p: power factor to generate fake data shape if dim of expected shape is "x**p", default to 1.
n: multiply factor to generate fake data shape if dim of expected shape is "x*n", default to 1.
any: specified size to generate fake data shape if dim of expected shape is "*", default to 1.
args_file: a JSON or YAML file to provide default values for `net_id`, `meta_file`, `config_file`,
`device`, `p`, `n`, `any`, and override pairs. so that the command line inputs can be simplified.
override: id-value pairs to override or add the corresponding config content.
e.g. ``--_meta#network_data_format#inputs#image#num_channels 3``.
"""
_args = _update_args(
args=args_file,
net_id=net_id,
meta_file=meta_file,
config_file=config_file,
device=device,
p=p,
n=n,
any=any,
**override,
)
_log_input_summary(tag="verify_net_in_out", args=_args)
config_file_, meta_file_, net_id_, device_, p_, n_, any_ = _pop_args(
_args,
"config_file",
"meta_file",
net_id="",
device="cuda:0" if is_available() else "cpu",
p=1,
n=1,
any=1)
parser = ConfigParser()
parser.read_config(f=config_file_)
parser.read_meta(f=meta_file_)
# the rest key-values in the _args are to override config content
for k, v in _args.items():
parser[k] = v
try:
key: str = net_id_ # mark the full id when KeyError
net = parser.get_parsed_content(key).to(device_)
key = "_meta_#network_data_format#inputs#image#num_channels"
input_channels = parser[key]
key = "_meta_#network_data_format#inputs#image#spatial_shape"
input_spatial_shape = tuple(parser[key])
key = "_meta_#network_data_format#inputs#image#dtype"
input_dtype = get_equivalent_dtype(parser[key], torch.Tensor)
key = "_meta_#network_data_format#outputs#pred#num_channels"
output_channels = parser[key]
key = "_meta_#network_data_format#outputs#pred#dtype"
output_dtype = get_equivalent_dtype(parser[key], torch.Tensor)
except KeyError as e:
raise KeyError(
f"Failed to verify due to missing expected key in the config: {key}."
) from e
net.eval()
with torch.no_grad():
spatial_shape = _get_fake_spatial_shape(input_spatial_shape,
p=p_,
n=n_,
any=any_) # type: ignore
test_data = torch.rand(*(1, input_channels, *spatial_shape),
dtype=input_dtype,
device=device_)
output = net(test_data)
if output.shape[1] != output_channels:
raise ValueError(
f"output channel number `{output.shape[1]}` doesn't match: `{output_channels}`."
)
if output.dtype != output_dtype:
raise ValueError(
f"dtype of output data `{output.dtype}` doesn't match: {output_dtype}."
)
logger.info("data shape of network is verified with no error.")
def load_bundle_config(bundle_path: str, *config_names, **load_kw_args) -> Any:
"""
Load the metadata and nominated configuration files from a MONAI bundle without loading the network itself.
This function will load the information from the bundle, which can be a directory or a zip file containing a
directory or a Torchscript bundle, and return the parser object with the information. This saves having to load
the model if only the information is wanted, and can work on any sort of bundle format.
Args:
bundle_path: path to the bundle directory or zip file
config_names: names of configuration files with extensions to load, should not be full paths but just name+ext
load_kw_args: keyword arguments to pass to the ConfigParser object when loading
Returns:
ConfigParser object containing the parsed information
"""
from monai.bundle.config_parser import ConfigParser # avoids circular import
parser = ConfigParser()
if not os.path.exists(bundle_path):
raise ValueError(f"Cannot find bundle file/directory '{bundle_path}'")
# bundle is a directory, read files directly
if os.path.isdir(bundle_path):
conf_data = []
parser.read_meta(f=os.path.join(bundle_path, "configs", "metadata.json"), **load_kw_args)
for cname in config_names:
cpath = os.path.join(bundle_path, "configs", cname)
if not os.path.exists(cpath):
raise ValueError(f"Cannot find config file '{cpath}'")
conf_data.append(cpath)
parser.read_config(f=conf_data, **load_kw_args)
else:
# bundle is a zip file which is either a zipped directory or a Torchscript archive
name, _ = os.path.splitext(os.path.basename(bundle_path))
archive = zipfile.ZipFile(bundle_path, "r")
all_files = archive.namelist()
zip_meta_name = f"{name}/configs/metadata.json"
if zip_meta_name in all_files:
prefix = f"{name}/configs/" # zipped directory location for files
else:
zip_meta_name = f"{name}/extra/metadata.json"
prefix = f"{name}/extra/" # Torchscript location for files
meta_json = json.loads(archive.read(zip_meta_name))
parser.read_meta(f=meta_json)
for cname in config_names:
full_cname = prefix + cname
if full_cname not in all_files:
raise ValueError(f"Cannot find config file '{full_cname}'")
ardata = archive.read(full_cname)
if full_cname.lower().endswith("json"):
cdata = json.loads(ardata, **load_kw_args)
elif full_cname.lower().endswith(("yaml", "yml")):
cdata = yaml.safe_load(ardata, **load_kw_args)
parser.read_config(f=cdata)
return parser