Repository: Project-MONAI/MONAI
File: monai/bundle/scripts.py (Python)

Code before:

    # the rest key-values in the _args are to override config content
    for k, v in _args.items():
        parser[k] = v

    # The convert_to_onnx must have an `inputs` as input, no matter what the `use_trace` is.
    # If the `input_shape` is not provided, will try to parse it from the parser to generate a random `inputs`.
    if not input_shape_:
        input_shape_ = _get_fake_input_shape(parser=parser)

    inputs_ = [torch.rand(input_shape_)]
    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:
        ckpt = torch.load(ckpt_file_)
        copy_model_state(dst=net, src=ckpt if key_in_ckpt_ == "" else ckpt[key_in_ckpt_])

    converter_kwargs_.update({"inputs": inputs_, "use_trace": use_trace_})
    onnx_model = convert_to_onnx(model=net, **converter_kwargs_)
    onnx.save(onnx_model, filepath_)


def ckpt_export(
    net_id: str | None = "network_def",
    filepath: PathLike | None = "models/model.ts",
    ckpt_file: str | None = "models/model.pt",
    meta_file: str | Sequence[str] | None = "configs/metadata.json",
    config_file: str | Sequence[str] | None = None,
    key_in_ckpt: str | None = None,
    use_trace: bool | None = None,
    input_shape: Sequence[int] | None = None,
    args_file: str | None = None,
    converter_kwargs: Mapping | None = None,
    **override: Any,
) -> None:
    """
    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.

Review comment (suggestion):
I didn't add `deprecated_arg_default` here since no breaking change will be introduced, what do you think?
cc @wyli @ericspod 

Code after:

    # the rest key-values in the _args are to override config content
    for k, v in _args.items():
        parser[k] = v

    # The convert_to_onnx must have an `inputs` as input, no matter what the `use_trace` is.
    # If the `input_shape` is not provided, will try to parse it from the parser to generate a random `inputs`.
    if not input_shape_:
        input_shape_ = _get_fake_input_shape(parser=parser)

    inputs_ = [torch.rand(input_shape_)]
    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:
        ckpt = torch.load(ckpt_file_)
        copy_model_state(dst=net, src=ckpt if key_in_ckpt_ == "" else ckpt[key_in_ckpt_])

    converter_kwargs_.update({"inputs": inputs_, "use_trace": use_trace_})
    onnx_model = convert_to_onnx(model=net, **converter_kwargs_)
    onnx.save(onnx_model, filepath_)


def ckpt_export(
    net_id: str | None = None,
    filepath: PathLike | None = None,
    ckpt_file: str | None = None,
    meta_file: str | Sequence[str] | None = None,
    config_file: str | Sequence[str] | None = None,
    key_in_ckpt: str | None = None,
    use_trace: bool | None = None,
    input_shape: Sequence[int] | None = None,
    args_file: str | None = None,
    converter_kwargs: Mapping | None = None,
    **override: Any,
) -> None:
    """
    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`.
            Default to "network_def".
        filepath: filepath to export, if filename has no extension it becomes `.ts`.
            Default to "models/model.ts" under "os.getcwd()" if `bundle_root` is not specified.