tinynn/converter/utils/tensorrt.py

import os
import re

from typing import Optional

import onnx
import onnxruntime as ort
import numpy as np
import pycuda.autoinit  # noqa: F401
import pycuda.driver as cuda
import tensorrt as trt

from contextlib import contextmanager


# A logger with the specific log level
# Available levels: WARNING, ERROR, VERBOSE, INFO
TRT_BUILD_LOGGER = trt.Logger(trt.Logger.WARNING)
TRT_EVAL_LOGGER = trt.Logger(trt.Logger.ERROR)

# Constant batch, currently it cannot be disabled
EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)


class MyLogger(trt.ILogger):
    def __init__(self, inner):
        trt.ILogger.__init__(self)
        self.info_log = None
        self.inner = inner

    def log(self, severity, msg):
        if msg.startswith('Engine Layer Information:'):
            self.info_log = msg
        self.inner.log(severity, msg)


class HostDeviceMem(object):
    def __init__(self, host_mem, device_mem):
        self.host = host_mem
        self.device = device_mem

    def __str__(self):
        return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device)

    def __repr__(self):
        return self.__str__()


def allocate_buffers(engine, dynamic_shapes_for_build, dynamc_shapes_for_eval):
    """Allocates all buffers required for an engine, i.e. host/device inputs/outputs."""
    inputs = []
    outputs = []
    bindings = []
    names_map_i = {}
    names_map_o = {}
    stream = cuda.Stream()
    for idx, binding in enumerate(engine):
        if binding in dynamc_shapes_for_eval:
            size = trt.volume(dynamc_shapes_for_eval[binding])
        elif binding in dynamic_shapes_for_build:
            size = trt.volume(dynamic_shapes_for_build[binding][1])
        else:
            size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size

        dtype = trt.nptype(engine.get_binding_dtype(binding))
        # Allocate host and device buffers
        host_mem = cuda.pagelocked_empty(size, dtype)
        device_mem = cuda.mem_alloc(host_mem.nbytes)
        # Append the device buffer to device bindings.
        bindings.append(int(device_mem))
        # Input/Output node name
        if engine.binding_is_input(binding):
            names_map_i[binding] = len(inputs)
            inputs.append(HostDeviceMem(host_mem, device_mem))
            print('input:', engine.get_binding_shape(binding), dtype, binding)
        else:
            names_map_o[binding] = len(outputs)
            outputs.append(HostDeviceMem(host_mem, device_mem))
            print('output:', engine.get_binding_shape(binding), dtype, binding)
    return inputs, outputs, bindings, stream, names_map_i, names_map_o


def do_inference(context, bindings, inputs, outputs, stream, batch_size=1):
    """This function is generalized for multiple inputs/outputs.
    inputs and outputs are expected to be lists of HostDeviceMem objects."""
    # Transfer input data to the GPU.
    [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]
    # Run inference.
    context.execute_async(batch_size=batch_size, bindings=bindings, stream_handle=stream.handle)
    # Transfer predictions back from the GPU.
    [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]
    # Synchronize the stream
    stream.synchronize()
    # Return only the host outputs.
    return [out.host for out in outputs]


def GiB(val):
    """Numerical value for GiB"""
    return val * 1 << 30


def add_profile(builder, config, dynamic_shapes_for_build):
    if len(dynamic_shapes_for_build) > 0:
        profile = builder.create_optimization_profile()
        for inp, (min_shape, opt_shape, max_shape) in dynamic_shapes_for_build.items():
            profile.set_shape(inp, min_shape, opt_shape, max_shape)
        config.add_optimization_profile(profile)


@contextmanager
def build_engine(model_path, logger, build_with_fp16, build_with_int8, build_with_workspace, dynamic_shapes_for_build):
    """Build the TensorRT engine from a ONNX model"""
    with trt.Builder(logger) as builder, builder.create_network(EXPLICIT_BATCH) as network, trt.OnnxParser(
        network, logger
    ) as parser:
        with open(model_path, 'rb') as model:
            if not parser.parse(model.read()):
                for error in range(parser.num_errors):
                    print(parser.get_error(error))

        with builder.create_builder_config() as config:
            if build_with_fp16:
                config.flags = 1 << (int)(trt.BuilderFlag.FP16)
            if build_with_int8:
                config.flags = 1 << (int)(trt.BuilderFlag.INT8)
            config.max_workspace_size = GiB(build_with_workspace)

            add_profile(builder, config, dynamic_shapes_for_build)

            yield builder.build_engine(network, config)


def get_needed_outputs(lines):
    needed_outputs = set()
    dtype_start = None
    shape_start = None
    start_end_dict = {'[': ']', '(': ')'}
    dtype_end = None
    shape_end = None
    sep = None
    for line in lines:
        if line.startswith('Layer('):
            line = line.rstrip('\n')

            op = re.findall(r'Layer\((.*?)\)', line)[0]

            if '->' not in line:
                continue

            pos = line.find('-> ')
            outputs = line[pos + 3 :]

            if dtype_start is None:
                match = next(re.finditer(r'( *(\[|\())(Int32|Float|Half|Int8)(\(|\[)', outputs))
                dtype_start = match.group(1)
                shape_start = match.group(4)
                dtype_end = start_end_dict[dtype_start.lstrip()]
                shape_end = start_end_dict[shape_start]

            if sep is None:
                end_str = f'{shape_end}{dtype_end}'
                pos = outputs.find(end_str)
                assert pos > 0
                start = pos + len(shape_end) + len(dtype_end)
                if start != len(outputs):
                    end = outputs.find(' ', start)
                    sep = f'{end_str}{outputs[start:end]} '

            if sep is not None and sep in outputs:
                outputs = outputs.split(sep)
            else:
                outputs = [outputs]

            for output in outputs:
                if dtype_start not in output:
                    continue

                pos = output.rfind(dtype_start)
                output_name = output[:pos]

                if '[' in output_name or '(' in output_name or '+' in output_name:
                    continue

                if re.match('Reformatted .* to .*', output):
                    continue

                print('Observing', op, repr(output_name))
                needed_outputs.add(output_name)
    return needed_outputs


def add_outputs_for_onnx_model(needed_outputs, onnx_path, new_onnx_path):
    model = onnx.load(onnx_path)

    orig_outputs = set()
    for node in model.graph.output:
        orig_outputs.add(node.name)
    print()

    for node in model.graph.node:
        for output in node.output:
            if output in needed_outputs and output not in orig_outputs:
                model.graph.output.extend([onnx.ValueInfoProto(name=output)])

                print('Added output:', output)

    inferred_model = onnx.shape_inference.infer_shapes(model)

    onnx.save_model(inferred_model, new_onnx_path)
    print('Modified model saved at', new_onnx_path)


def compare_onnx_tensorrt(
    onnx_path: str,
    build_with_fp16: bool,
    build_with_int8: bool,
    build_with_workspace: int = 4,
    dynamic_shapes_for_build: Optional[dict] = None,
    dynamc_shapes_for_eval: Optional[dict] = None,
    input_path_mapping: Optional[dict] = None,
):
    if dynamic_shapes_for_build is None:
        dynamic_shapes_for_build = {}

    if dynamc_shapes_for_eval is None:
        dynamc_shapes_for_eval = {}

    if input_path_mapping is None:
        input_path_mapping = {}

    onnx_fn, onnx_ext = os.path.splitext(onnx_path)
    new_onnx_path = f'{onnx_fn}_with_outputs{onnx_ext}'
    new_trt_path = f'{onnx_fn}_with_outputs{onnx_ext}'

    print('Building TensorRT engine with', onnx_path)
    logger = MyLogger(TRT_BUILD_LOGGER)
    with build_engine(
        onnx_path, logger, build_with_fp16, build_with_int8, build_with_workspace, dynamic_shapes_for_build
    ) as engine:
        pass

    assert logger.info_log is not None, "Engine layer information is missing"
    lines = logger.info_log.splitlines()
    needed_outputs = get_needed_outputs(lines)
    add_outputs_for_onnx_model(needed_outputs, onnx_path, new_onnx_path)

    with build_engine(
        new_onnx_path, TRT_BUILD_LOGGER, build_with_fp16, build_with_int8, build_with_workspace
    ) as engine:
        with open(new_trt_path, 'wb') as f:
            f.write(bytearray(engine.serialize()))

    runtime = trt.Runtime(TRT_EVAL_LOGGER)
    with open(new_trt_path, 'rb') as f:
        engine_bytes = f.read()
        engine = runtime.deserialize_cuda_engine(engine_bytes)

    print('=' * 60)
    print('input output tensors:')
    inputs, outputs, bindings, stream, names_map_i, names_map_o = allocate_buffers(
        engine, dynamic_shapes_for_build, dynamc_shapes_for_eval
    )
    print('=' * 60)

    input_data = {}

    with engine.create_execution_context() as context:
        if True:
            for binding in names_map_i:
                dtype = trt.nptype(engine.get_binding_dtype(binding))

                if binding in input_path_mapping:
                    assert (
                        binding in dynamc_shapes_for_eval
                    ), "input_path_mapping and dynamc_shapes_for_eval should be specified together"

                    shape = dynamc_shapes_for_eval[binding]
                    data = np.fromfile(input_path_mapping[binding], dtype='uint8')
                    data = np.reshape(data.view(dtype), shape)
                    input_data[binding] = data

                    context.set_binding_shape(engine.get_binding_index(binding), shape)
                else:
                    if binding in dynamc_shapes_for_eval:
                        shape = dynamc_shapes_for_eval[binding]
                        context.set_binding_shape(engine.get_binding_index(binding), shape)
                    else:
                        shape = engine.get_binding_shape(binding)
                        if -1 in shape:
                            shape = dynamic_shapes_for_build[binding][1]
                            context.set_binding_shape(engine.get_binding_index(binding), shape)

                    data = np.random.random(shape).astype(dtype)
                    input_data[binding] = data

                np.copyto(inputs[names_map_i[binding]].host, data.ravel())

            output = do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)

            output_keys = list(names_map_o)

    opts = ort.SessionOptions()
    opts.intra_op_num_threads = 1
    opts.inter_op_num_threads = 1
    sess = ort.InferenceSession(new_onnx_path, providers=['CUDAExecutionProvider'], sess_options=opts)

    res = sess.run(output_keys, input_data)

    for i, binding in enumerate(output_keys):
        onnx_val = res[i]

        if onnx_val is None:
            continue

        dtype = trt.nptype(engine.get_binding_dtype(binding))

        if dtype != onnx_val.dtype:
            trt_val = np.reshape(output[names_map_o[binding]].view(dtype).astype(onnx_val.dtype), onnx_val.shape)
        else:
            trt_val = np.reshape(output[names_map_o[binding]].view(onnx_val.dtype), onnx_val.shape)

        is_aligned = np.allclose(onnx_val, trt_val)

        if is_aligned:
            print(binding, 'matches:', is_aligned)
            continue

        onnx_val_ravel = onnx_val.ravel()
        trt_val_ravel = trt_val.ravel()

        cross_sim = np.dot(onnx_val_ravel, trt_val_ravel) / (
            np.linalg.norm(onnx_val_ravel) * np.linalg.norm(trt_val_ravel)
        )
        is_aligned = is_aligned or cross_sim > 0.999
        print(binding, 'matches:', is_aligned, 'cross_sim =', cross_sim)

        if not is_aligned:
            print('Top 10 values with maximum differences')
            max_diff_indices = np.argsort(np.abs(onnx_val - trt_val).ravel())[::-1][:10]
            print('TensorRT:')
            print(trt_val_ravel[max_diff_indices])
            print('ONNX:')
            print(onnx_val_ravel[max_diff_indices])

        print('-' * 60)