def test_script_save(net, *inputs, device=None, rtol=1e-4, atol=0.0):
"""
Test the ability to save `net` as a Torchscript object, reload it, and apply inference. The value `inputs` is
forward-passed through the original and loaded copy of the network and their results returned.
The forward pass for both is done without gradient accumulation.
The test will be performed with CUDA if available, else CPU.
"""
# TODO: would be nice to use GPU if available, but it currently causes CI failures.
device = "cpu"
with tempfile.TemporaryDirectory() as tempdir:
convert_to_torchscript(
model=net,
filename_or_obj=os.path.join(tempdir, "model.ts"),
verify=True,
inputs=inputs,
device=device,
rtol=rtol,
atol=atol,
)
def test_value(self):
model = UNet(spatial_dims=2,
in_channels=1,
out_channels=3,
channels=(16, 32, 64),
strides=(2, 2),
num_res_units=0)
with tempfile.TemporaryDirectory() as tempdir:
torchscript_model = convert_to_torchscript(
model=model,
filename_or_obj=os.path.join(tempdir, "model.ts"),
extra_files={"foo.txt": b"bar"},
verify=True,
inputs=[torch.randn((16, 1, 32, 32), requires_grad=False)],
device="cuda" if torch.cuda.is_available() else "cpu",
rtol=1e-3,
atol=1e-4,
)
self.assertTrue(isinstance(torchscript_model, torch.nn.Module))
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_}.")