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

Code before:
            --ckpt_file <checkpoint path> --input_shape <input shape> --dynamic_batchsize <batch range> ...

    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 the TensorRT based torchscript model and extract network
            information, the saved key in the 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.
        precision: the weight precision of the converted TensorRT engine based torchscript models. Should be 'fp32' or 'fp16'.
        input_shape: the input shape that is used to convert the 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.
        use_trace: whether using `torch.jit.trace` to convert the pytorch model to torchscript model and then convert to
            a TensorRT engine based torchscript model.
        dynamic_batchsize: a sequence with three elements to define the batch size range of the input for the model to be
            converted. Should be a sequence like [MIN_BATCH, OPT_BATCH, MAX_BATCH]. After converted, the batchsize of
            model input should between `MIN_BATCH` and `MAX_BATCH` and the `OPT_BATCH` is the best performance batchsize
            that the TensorRT tries to fit. The `OPT_BATCH` should be the most frequently used input batchsize in
            the application.
        device: the target GPU index to convert and verify the model.
        use_onnx: whether to use the ONNX-TensorRT way to export the TensorRT engine-based torchscript model.
        onnx_input_names: optional input names of the ONNX model.
        onnx_output_names: optional output names of the ONNX model.
        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_trt`, 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_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,
        precision=precision,
        input_shape=input_shape,
        use_trace=use_trace,
        dynamic_batchsize=dynamic_batchsize,
        device=device,
        use_onnx=use_onnx,
        onnx_input_names=onnx_input_names,
        onnx_output_names=onnx_output_names,

Review comment (suggestion):
Suggest to add a simple description of the default value, it seems "input_0" is used.
Same suggestion on `onnx_output_names `

Code after:
    When exporting through the first way, some models suffer from the slowdown problem, since Torch-TensorRT
    may only convert a little part of the PyTorch model to the TensorRT engine. However when exporting through
    the second way, some Python data structures like `dict` are not supported. And some TorchScript models are
    not supported by the ONNX if exported through `torch.jit.script`.

    Typical usage examples:

    .. code-block:: bash

        python -m monai.bundle trt_export --net_id <network definition> --filepath <export path> \
            --ckpt_file <checkpoint path> --input_shape <input shape> --dynamic_batchsize <batch range> ...

    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 the TensorRT based TorchScript model and extract network
            information, the saved key in the 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.
        precision: the weight precision of the converted TensorRT engine based TorchScript models. Should be 'fp32' or 'fp16'.
        input_shape: the input shape that is used to convert the 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.
        use_trace: whether using `torch.jit.trace` to convert the PyTorch model to a TorchScript model and then convert to
            a TensorRT engine based TorchScript model or an ONNX model (if `use_onnx` is True).
        dynamic_batchsize: a sequence with three elements to define the batch size range of the input for the model to be
            converted. Should be a sequence like [MIN_BATCH, OPT_BATCH, MAX_BATCH]. After converted, the batchsize of
            model input should between `MIN_BATCH` and `MAX_BATCH` and the `OPT_BATCH` is the best performance batchsize
            that the TensorRT tries to fit. The `OPT_BATCH` should be the most frequently used input batchsize in
            the application.
        device: the target GPU index to convert and verify the model.
        use_onnx: whether using the ONNX-TensorRT way to export the TensorRT engine-based TorchScript model.
        onnx_input_names: optional input names of the ONNX model. This arg is only useful when `use_onnx` is True. Should be
            a sequence like `['input_0', 'input_1', ..., 'input_N']` where N equals to the number of the model inputs. If not
            given, will use `['input_0']`, which supposes the model only has one input.
        onnx_output_names: optional output names of the ONNX model. This arg is only useful when `use_onnx` is True. Should be
            a sequence like `['output_0', 'output_1', ..., 'output_N']` where N equals to the number of the model outputs. If
            not given, will use `['output_0']`, which supposes the model only has one output.
        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_trt`, 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_args(
        args=args_file,