Setting Mixed-Precision on weights, and not only computation type for TensorRT-8.6.3

[bernardrb]

Description

In the code I'm basically trying to see if I can enable int8, and thereafter set the layer precision back to fp16. I want to understand how I can use mixed-precision in tensorrt.

I set computation type to FP16 using: self.network.get_layer(i).precision = trt.DataType.HALF

I try to set weights using:
for j in range(layer.num_outputs: layer.set_output_type(j, trt.DataType.FLOAT)

Using EngineInspector, I still see that all weights are in int8. As seen by the below log. I also can't find LayerType Reformat in the python tensorrt api, and I don't implement any plugin layers.

{
"Name": "Reformatting CopyNode for Input Tensor 0 to /image_encoder/backbone/stages.0/op_list.0/conv/Conv",
"LayerType": "Reformat",
"Inputs": [
{
"Name": "input_image",
"Location": "Device",
"Dimensions": [1,3,1024,1024],
"Format/Datatype": "Row major linear FP32"
}],
"Outputs": [
{
"Name": "Reformatted Input Tensor 0 to /image_encoder/backbone/stages.0/op_list.0/conv/Conv",
"Location": "Device",
"Dimensions": [1,3,1024,1024],
"Format/Datatype": "Row major Int8 format"
}],
"ParameterType": "Reformat",
"Origin": "REFORMAT",
"TacticValue": "0x0000000000000000",
"StreamId": 0,
"Metadata": ""
}
{
"Name": "/image_encoder/backbone/stages.0/op_list.0/conv/Conv",
"LayerType": "CaskConvolution",
"Inputs": [
{
"Name": "Reformatted Input Tensor 0 to /image_encoder/backbone/stages.0/op_list.0/conv/Conv",
"Location": "Device",
"Dimensions": [1,3,1024,1024],
"Format/Datatype": "Row major Int8 format"
}],
"Outputs": [
{
"Name": "/image_encoder/backbone/stages.0/op_list.0/conv/Conv_output_0",
"Location": "Device",
"Dimensions": [1,32,512,512],
"Format/Datatype": "Thirty-two wide channel vectorized row major Int8 format"
}],
"ParameterType": "Convolution",
"Kernel": [3,3],
"PaddingMode": "kEXPLICIT_ROUND_DOWN",
"PrePadding": [1,1],
"PostPadding": [1,1],
"Stride": [2,2],
"Dilation": [1,1],
"OutMaps": 32,
"Groups": 1,
"Weights": {"Type": "Int8", "Count": 864},
"Bias": {"Type": "Float", "Count": 32},
"HasSparseWeights": 0,
"HasDynamicFilter": 0,
"HasDynamicBias": 0,
"HasResidual": 0,
"ConvXAsActInputIdx": -1,
"BiasAsActInputIdx": -1,
"ResAsActInputIdx": -1,
"Activation": "NONE",
"HasBias": 1,
"HasReLU": 0,
"TacticName": "ampere_first_layer_filter3x3_imma_fwd",
"TacticValue": "0x9ae0c0d2fb3a01e5",
"StreamId": 0,
"Metadata": "[ONNX Layer: /image_encoder/backbone/stages.0/op_list.0/conv/Conv]"
}
{"Layers": [{
"Name": "Reformatting CopyNode for Input Tensor 0 to /image_encoder/backbone/stages.0/op_list.0/conv/Conv",
"LayerType": "Reformat",
"Inputs": [
{
"Name": "input_image",
"Location": "Device",
"Dimensions": [1,3,1024,1024],
"Format/Datatype": "Row major linear FP32"
}],
"Outputs": [
{
"Name": "Reformatted Input Tensor 0 to /image_encoder/backbone/stages.0/op_list.0/conv/Conv",
"Location": "Device",
"Dimensions": [1,3,1024,1024],
"Format/Datatype": "Row major Int8 format"
}],
"ParameterType": "Reformat",
"Origin": "REFORMAT",
"TacticValue": "0x0000000000000000",
"StreamId": 0,
"Metadata": ""
},{
"Name": "/image_encoder/backbone/stages.0/op_list.0/conv/Conv",
"LayerType": "CaskConvolution",
"Inputs": [
{
"Name": "Reformatted Input Tensor 0 to /image_encoder/backbone/stages.0/op_list.0/conv/Conv",
"Location": "Device",
"Dimensions": [1,3,1024,1024],
"Format/Datatype": "Row major Int8 format"
}],
"Outputs": [
{
"Name": "/image_encoder/backbone/stages.0/op_list.0/conv/Conv_output_0",
"Location": "Device",
"Dimensions": [1,32,512,512],
"Format/Datatype": "Thirty-two wide channel vectorized row major Int8 format"
}],
"ParameterType": "Convolution",
"Kernel": [3,3],
"PaddingMode": "kEXPLICIT_ROUND_DOWN",
"PrePadding": [1,1],
"PostPadding": [1,1],
"Stride": [2,2],
"Dilation": [1,1],
"OutMaps": 32,
"Groups": 1,
"Weights": {"Type": "Int8", "Count": 864},
"Bias": {"Type": "Float", "Count": 32},
"HasSparseWeights": 0,
"HasDynamicFilter": 0,
"HasDynamicBias": 0,
"HasResidual": 0,
"ConvXAsActInputIdx": -1,
"BiasAsActInputIdx": -1,
"ResAsActInputIdx": -1,
"Activation": "NONE",
"HasBias": 1,
"HasReLU": 0,
"TacticName": "ampere_first_layer_filter3x3_imma_fwd",
"TacticValue": "0x9ae0c0d2fb3a01e5",
"StreamId": 0,
"Metadata": "[ONNX Layer: /image_encoder/backbone/stages.0/op_list.0/conv/Conv]"
}

def set_mixed_precision(self):

    """
    Experimental precision mode.
    Enable mixed-precision mode. When set, the layers defined here will be forced to FP16 to maximize
    INT8 inference accuracy, while having minimal impact on latency.
    """

    self.config.set_flag(trt.BuilderFlag.OBEY_PRECISION_CONSTRAINTS)

    # All convolution operations in the first four blocks of the graph are pinned to FP16.
    # These layers have been manually chosen as they give a good middle-point between int8 and fp16
    # accuracy in COCO, while maintining almost the same latency as a normal int8 engine.
    # To experiment with other datasets, or a different balance between accuracy/latency, you may
    # add or remove blocks.
    for i in range(self.network.num_layers):
        layer = self.network.get_layer(i)
        if layer.precision != trt.DataType.FLOAT and layer.type != trt.LayerType.SHAPE and layer.type != trt.LayerType.NORMALIZATION and "norm/" not in layer.name: 
            self.network.get_layer(i).precision = trt.DataType.HALF
            for j in range(layer.num_outputs):
                layer.set_output_type(j, trt.DataType.FLOAT)
        logger.info("Mixed-Precision Layer {} set to HALF STRICT data type".format(layer.name))
        if "norm/" in layer.name or layer.type == trt.LayerType.NORMALIZATION:
            self.network.get_layer(i).precision = trt.DataType.FLOAT
            logger.info("Mixed-Precision Layer {} set to FLOAT STRICT data type".format(layer.name))

def create_engine(self, engine_path, precision, img_size, calib_input=None, calib_cache=None, calib_num_images=5000,
                  calib_batch_size=8):
    """
    Build the TensorRT engine and serialize it to disk.
    :param engine_path: The path where to serialize the engine to.
    :param precision: The datatype to use for the engine, either 'fp32', 'fp16', 'int8', or 'mixed'.
    :param calib_input: The path to a directory holding the calibration images.
    :param calib_cache: The path where to write the calibration cache to, or if it already exists, load it from.
    :param calib_num_images: The maximum number of images to use for calibration.
    :param calib_batch_size: The batch size to use for the calibration process.
    """
    engine_path = os.path.realpath(engine_path)
    engine_dir = os.path.dirname(engine_path)
    os.makedirs(engine_dir, exist_ok=True)
    logger.info("Building {} Engine in {}".format(precision, engine_path))

    inputs = [self.network.get_input(i) for i in range(self.network.num_inputs)]

    if precision in ["fp16", "int8", "mixed"]:
        if not self.builder.platform_has_fast_fp16:
            log.warning("FP16 is not supported natively on this platform/device")
        self.config.set_flag(trt.BuilderFlag.FP16)
    if precision in ["int8", "mixed"]:
        if not self.builder.platform_has_fast_int8:
            log.warning("INT8 is not supported natively on this platform/device")
        self.config.set_flag(trt.BuilderFlag.INT8)
        self.config.int8_calibrator = EngineCalibrator(calib_cache)
        logger.debug("int8_calibrator.get_algorithm(): {}".format(self.config.int8_calibrator.get_algorithm()))
        if calib_cache is None or not os.path.exists(calib_cache):
            calib_shape = [calib_batch_size] + list(inputs[0].shape[1:])
            calib_dtype = trt.nptype(inputs[0].dtype)
            logger.debug("Calibration shape: {}".format(calib_shape))
            logger.debug("Calibration dtype: {}".format(calib_dtype))
            self.config.int8_calibrator.set_image_batcher(
                ImageBatcher(calib_input, calib_shape, calib_dtype, img_size=img_size, max_num_images=calib_num_images,
                             exact_batches=True, shuffle_files=True))
    self.config.profiling_verbosity  = trt.ProfilingVerbosity.DETAILED

    engine_bytes = None
    try:
        engine_bytes = self.builder.build_serialized_network(self.network, self.config)
    except AttributeError:
        engine = self.builder.build_engine(self.network, self.config)
        engine_bytes = engine.serialize()
        del engine
    assert engine_bytes
    with open(engine_path, "wb") as f:
        logger.info("Serializing engine to file: {:}".format(engine_path))
        f.write(engine_bytes)

def main(args):
builder = EngineBuilder(args.verbose, args.workspace)
builder.create_network(args.onnx, args.batch_size, args.dynamic_batch_size)
if args.precision == "mixed":
builder.set_mixed_precision()
builder.create_engine(args.engine, args.precision, args.img_size, args.calib_input, args.calib_cache, args.calib_num_images,
args.calib_batch_size)

Environment

TensorRT Version: 8.6.3

NVIDIA GPU: RTX 3090

NVIDIA Driver Version:

CUDA Version: 12.0

CUDNN Version:

Baremetal or Container (if so, version): 24.02-py3

[zerollzeng]

Please refer to https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#layer-level-control

[zerollzeng]

Basically compute precision decide input and compute precision, while output precision decides output precision.

[bernardrb]

In https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#layer-level-control,

By default, TensorRT chooses such an implementation only if it results in a higher-performance network. If another implementation is faster, TensorRT uses it and issues a warning. You can override this behavior by preferring the type constraints in the builder configuration. config.set_flag(trt.BuilderFlag.OBEY_PRECISION_CONSTRAINTS);

However, in https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#explicit-implicit-quantization

In explicitly quantized networks, the scaling operations to transform between the quantized and unquantized values are represented explicitly by IQuantizeLayer and IDequantizeLayer nodes in the graph - these will henceforth be referred to as Q/DQ nodes. By contrast with implicit quantization, the explicit form specifies exactly where conversion to and from INT8 is performed, and the optimizer will perform only precision conversions that are dictated by the semantics of the model.

What I don't get is that setting TensorRT to obey precisions forces it to use the specified precision. However, later it says TensorRT uses implicit quantization, and therefore one does not have control of layer-wise precision. It is then concluded that explicit quantization is only achieved by adding quantization layers in e.g. PyTorch, and then exporting to ONNX.

Is setting TensorRT to obey precision constraints the same thing as adding quantization layers in PyTorch (i.e would be categorized as explicit quantization)? If they are the same, are there hidden trade-offs or issues with either approach?

Please correct any misunderstandings!

[zerollzeng]

1. If a node has explicit quantization(Q/DQ) then you cannot set the layer precision for it since it's been fixed.
2. for implicit quantization, you can fallback those int8 layers back to FP16/FP32.

You may want to check https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#strong-vs-weak-typing, I know this a bit of confusing due to historical reason...

[zerollzeng]

The general rule is: only set layer precision for specific reason, e.g. int8 accuracy is not good. otherwise let TensorRT do the optimization.

[zerollzeng]

I only use this feature to fallback some layers back to FP16/FP32 to keep accuracy quickly, e.g. I have a PTQ engine that don't have good accuracy, I may fallback some layers back to FP16/FP32 to optimize, although it will lose some performance, the best solution in this case is finetune the model with QAT so that we can have "user controlled" accuracy.