Repository: Project-MONAI/MONAI
File: monai/networks/utils.py (Python)

Code before:
    # default to an empty list to fit the `torch_tensorrt.ts.embed_engine_in_new_module` function.
    input_names = [] if not input_names else input_names
    output_names = [] if not output_names else output_names

    # set up the TensorRT builder
    torch_tensorrt.set_device(device)
    logger = trt.Logger(trt.Logger.WARNING)
    builder = trt.Builder(logger)
    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    profile = builder.create_optimization_profile()
    if input_names:
        profile.set_shape(input_names[0], *input_shapes)

    # parse the ONNX model
    parser = trt.OnnxParser(network, logger)
    success = parser.parse(onnx_model.SerializeToString())
    for idx in range(parser.num_errors):
        print(parser.get_error(idx))

    if not success:
        raise Exception("Cannot export onnx model to trt due to upper error.")

    # set up the conversion configuration
    config = builder.create_builder_config()
    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 31)
    config.add_optimization_profile(profile)
    if precision == "fp16":
        config.set_flag(trt.BuilderFlag.FP16)
    serialized_engine = builder.build_serialized_network(network, config)
    f = io.BytesIO()
    f.write(serialized_engine)

    # wrap the serialized TensorRT engine back to a TorchScript module.
    trt_model = torch_tensorrt.ts.embed_engine_in_new_module(
        f.getvalue(), torch.device(f"cuda:{device}"), input_names, output_names
    )
    return trt_model


def convert_to_trt(
    model: nn.Module,
    precision: str,
    input_shape: Sequence[int],
    dynamic_batchsize: Sequence[int] | None = None,
    use_trace: bool = False,
    filename_or_obj: Any | None = None,
    verify: bool = False,
    device: int | None = None,
    use_onnx: bool | None = False,
    onnx_input_names: Sequence[str] | None = ("input_0",),
    onnx_output_names: Sequence[str] | None = ("output_0",),

Review comment (suggestion):
I think this logic can put before:
```
    for idx in range(parser.num_errors):
        print(parser.get_error(idx))
```

Code after:
        output_names: optional output names of the ONNX model. Should be a sequence like
            `['output_0', 'output_1', ..., 'output_N']` where N equals to the number of
            the model outputs.

    """
    trt, _ = optional_import("tensorrt", "8.5.3")
    torch_tensorrt, _ = optional_import("torch_tensorrt", "1.4.0")

    input_shapes = (min_shape, opt_shape, max_shape)
    # default to an empty list to fit the `torch_tensorrt.ts.embed_engine_in_new_module` function.
    input_names = [] if not input_names else input_names
    output_names = [] if not output_names else output_names

    # set up the TensorRT builder
    torch_tensorrt.set_device(device)
    logger = trt.Logger(trt.Logger.WARNING)
    builder = trt.Builder(logger)
    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    profile = builder.create_optimization_profile()
    if input_names:
        profile.set_shape(input_names[0], *input_shapes)

    # parse the ONNX model
    parser = trt.OnnxParser(network, logger)
    success = parser.parse(onnx_model.SerializeToString())
    if not success:
        parser_error_message = ""
        for idx in range(parser.num_errors):
            parser_error_message += parser.get_error(idx).desc() + "\n"
        raise Exception(f"TensorRT cannot parse the ONNX model, due to:\n{parser_error_message}")

    # set up the conversion configuration
    config = builder.create_builder_config()
    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 31)
    config.add_optimization_profile(profile)
    if precision == "fp16":
        config.set_flag(trt.BuilderFlag.FP16)
    serialized_engine = builder.build_serialized_network(network, config)
    f = io.BytesIO()
    f.write(serialized_engine)

    # wrap the serialized TensorRT engine back to a TorchScript module.
    trt_model = torch_tensorrt.ts.embed_engine_in_new_module(
        f.getvalue(), torch.device(f"cuda:{device}"), input_names, output_names
    )
    return trt_model


def convert_to_trt(
    model: nn.Module,
    precision: str,