Function build_trt_engine

tensorrt_llm/tools/multimodal_builder.py:178–355

Build TensorRT engine from ONNX model. Args: model_params (dict): Optional model specific parameters, e.g.: - qwen2_vl_dim (int): Dimension for Qwen2-VL model - min_hw_dims (int): Minimum HW dimensions - max_hw_dims (int): Maximum HW dimensions

(model_type,
                     input_sizes,
                     onnx_dir,
                     engine_dir,
                     max_batch_size,
                     dtype=torch.float16,
                     model_params=None,
                     onnx_name='model.onnx',
                     engine_name='model.engine',
                     delete_onnx=True,
                     logger=trt.Logger(trt.Logger.INFO))

def build_trt_engine(model_type,
                     input_sizes,
                     onnx_dir,
                     engine_dir,
                     max_batch_size,
                     dtype=torch.float16,
                     model_params=None,
                     onnx_name='model.onnx',
                     engine_name='model.engine',
                     delete_onnx=True,
                     logger=trt.Logger(trt.Logger.INFO)):
    """Build TensorRT engine from ONNX model.

    Args:
        model_params (dict): Optional model specific parameters, e.g.:
            - qwen2_vl_dim (int): Dimension for Qwen2-VL model
            - min_hw_dims (int): Minimum HW dimensions
            - max_hw_dims (int): Maximum HW dimensions
            - num_frames (int): Number of frames for video models
    """
    model_params = model_params or {}
    onnx_file = f'{onnx_dir}/{onnx_name}'
    engine_file = f'{engine_dir}/{engine_name}'
    config_file = f'{engine_dir}/config.json'
    logger.log(trt.Logger.INFO, f"Building TRT engine to {engine_file}")

    builder = trt.Builder(logger)
    network = builder.create_network(
        1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    profile = builder.create_optimization_profile()

    config_args = {
        "precision": torch_dtype_to_str(dtype),
        "model_type": model_type,
        "strongly_typed": False,
        "max_batch_size": max_batch_size,
        "model_name": "multiModal"
    }

    if "num_frames" in model_params:
        config_args["num_frames"] = model_params["num_frames"]

    config_wrapper = Builder().create_builder_config(**config_args)
    config = config_wrapper.trt_builder_config

    parser = trt.OnnxParser(network, logger)

    with open(onnx_file, 'rb') as model:
        if not parser.parse(model.read(), os.path.abspath(onnx_file)):
            logger.log(trt.Logger.ERROR, "Failed parsing %s" % onnx_file)
            for error in range(parser.num_errors):
                logger.log(trt.Logger.ERROR, parser.get_error(error))
        logger.log(trt.Logger.INFO, "Succeeded parsing %s" % onnx_file)

    nBS = -1
    nMinBS = 1
    nOptBS = max(nMinBS, int(max_batch_size / 2))
    nMaxBS = max_batch_size