def test_save_net_with_metadata_ext(self):
"""Save a network without metadata to a file."""
m = torch.jit.script(TestModule())
with tempfile.TemporaryDirectory() as tempdir:
save_net_with_metadata(m, f"{tempdir}/test.zip")
self.assertTrue(os.path.isfile(f"{tempdir}/test.zip"))
def test_save_net_with_metadata_with_extra(self):
"""Save a network with simple metadata to a file."""
m = torch.jit.script(TestModule())
test_metadata = {"foo": [1, 2], "bar": "string"}
with tempfile.TemporaryDirectory() as tempdir:
save_net_with_metadata(m,
f"{tempdir}/test",
meta_values=test_metadata)
self.assertTrue(os.path.isfile(f"{tempdir}/test.ts"))
def test_load_net_with_metadata(self):
"""Save then load a network with no metadata or other extra files."""
m = torch.jit.script(TestModule())
with tempfile.TemporaryDirectory() as tempdir:
save_net_with_metadata(m, f"{tempdir}/test")
_, meta, extra_files = load_net_with_metadata(f"{tempdir}/test.ts")
del meta[JITMetadataKeys.TIMESTAMP.
value] # no way of knowing precisely what this value would be
self.assertEqual(meta, get_config_values())
self.assertEqual(extra_files, {})
def test_load_net_with_metadata_with_extra(self):
"""Save then load a network with basic metadata."""
m = torch.jit.script(TestModule())
test_metadata = {"foo": [1, 2], "bar": "string"}
with tempfile.TemporaryDirectory() as tempdir:
save_net_with_metadata(m,
f"{tempdir}/test",
meta_values=test_metadata)
_, meta, extra_files = load_net_with_metadata(f"{tempdir}/test.ts")
del meta[JITMetadataKeys.TIMESTAMP.
value] # no way of knowing precisely what this value would be
test_compare = get_config_values()
test_compare.update(test_metadata)
self.assertEqual(meta, test_compare)
self.assertEqual(extra_files, {})
def test_save_load_more_extra_files(self):
"""Save then load extra file data from a torchscript file."""
m = torch.jit.script(TestModule())
test_metadata = {"foo": [1, 2], "bar": "string"}
more_extra_files = {"test.txt": b"This is test data"}
with tempfile.TemporaryDirectory() as tempdir:
save_net_with_metadata(m,
f"{tempdir}/test",
meta_values=test_metadata,
more_extra_files=more_extra_files)
self.assertTrue(os.path.isfile(f"{tempdir}/test.ts"))
_, _, loaded_extra_files = load_net_with_metadata(
f"{tempdir}/test.ts", more_extra_files=("test.txt", ))
self.assertEqual(more_extra_files["test.txt"],
loaded_extra_files["test.txt"])
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_}.")