Function onnx_export

monai/bundle/scripts.py:1323–1434

Export the model checkpoint to an onnx model. Typical usage examples: .. code-block:: bash python -m monai.bundle onnx_export network --filepath --ckpt_file ... Args: net_id: ID name of the network component in the config, it must

(
    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,
)

def onnx_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 an onnx model.

    Typical usage examples:

    .. code-block:: bash

        python -m monai.bundle onnx_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 where the onnx model is saved to.
        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 that contains extract network information,
        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.
        use_trace: whether using `torch.jit.trace` to convert the pytorch model to torchscript model.
        input_shape: a shape used to generate the random input of the network, when converting the model to an
            onnx model. Should be a list like [N, C, H, W] or [N, C, H, W, D]. If not given, will try to parse from
            the `metadata` config.
        args_file: a JSON or YAML file to provide default values for all the parameters of this function, so that
            the command line inputs can be simplified.
        converter_kwargs: extra arguments that are needed by `convert_to_onnx`, except ones that already exist in the
            input parameters.
        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_kwargs(
        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,
        use_trace=use_trace,
        input_shape=input_shape,
        converter_kwargs=converter_kwargs,
        **override,
    )
    _log_input_summary(tag="onnx_export", args=_args)
    (
        filepath_,