ONNX与TensorRT的使用

ONNXRuntime/TensorRT

1、利用TensortRT加速Tensorflow模型

主体思想：Tensorflow->TensorRT（pb->uff）

​ 以TensorRT官方样例基于手写数字体MNIST数据集的Lenet5模型为例。

​ 首先，下载MNIST数据集。

​ 然后训练手写数字体识别模型Lenet5。

​ 随后转换Lenet5.pb模型为Lenet5.uff模型。

# 下载MNIST数据集
cd /data/mnist
python download_pgms.py 

# 开始训练Lenet5模型
cd /samples/python/end_to_end_tensorflow_mnist/
python model.py
# 可得到模型Lenet5.pb

# 转换Lenet5.pb模型为Lenet5.uff模型
convert-to-uff ./models/Lenet5.pb ./models/Lenet5.uff

最后利用TensorRT加载Lenet5.uff并加速推理。

# -*- coding:UTF-8 -*-
from random import randint
from PIL import Image
import numpy as np
import pycuda.driver as cuda
import pycuda.autoinit
import tensorrt as trt
import sys, os
# 导入/samples/python/common.py
sys.path.insert(1, os.path.join(sys.path[0], ".."))
import common

TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

class ModelData(object):
    MODEL_FILE = "lenet5.uff"  # 模型路径
    INPUT_NAME ="input_1"  # 输入层名称 （convert-to-uff时可显示）
    INPUT_SHAPE = (1, 28, 28)  # 输入尺寸
    OUTPUT_NAME = "dense_1/Softmax"  # 输出层名称（convert-to-uff时可显示）

def build_engine(model_file):
    # For more information on TRT basics, refer to the introductory samples.
    with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.UffParser() as parser:
        builder.max_workspace_size = common.GiB(1)
        # Parse the Uff Network
        parser.register_input(ModelData.INPUT_NAME, ModelData.INPUT_SHAPE)
        parser.register_output(ModelData.OUTPUT_NAME)
        parser.parse(model_file, network)
        # Build and return an engine.
        return builder.build_cuda_engine(network)

# Loads a test case into the provided pagelocked_buffer.
def load_normalized_test_case(data_paths, pagelocked_buffer, case_num=randint(0, 9)):
    test_case_path = os.path.join(data_paths, str(case_num) + ".pgm")
    # Flatten the image into a 1D array, normalize, and copy to pagelocked memory.
    img = np.array(Image.open(test_case_path)).ravel()
    np.copyto(pagelocked_buffer, 1.0 - img / 255.0)
    return case_num

def main():
    data_paths = './data'
    model_path = os.environ.get("MODEL_PATH") or os.path.join(os.path.dirname(__file__), "models")
    model_file = os.path.join(model_path, ModelData.MODEL_FILE)

    with build_engine(model_file) as engine:
        inputs, outputs, bindings, stream = common.allocate_buffers(engine)
        with engine.create_execution_context() as context:
            case_num = load_normalized_test_case(data_paths, pagelocked_buffer=inputs[0].host)
            [output] = common.do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)
            pred = np.argmax(output)
            print("Test Case: " + str(case_num))
            print("Prediction: " + str(pred))

if __name__ == '__main__':
    main()

2、利用TensorRT加速Keras模型

主体思想：Keras->ONNX->TensorRT（h5->onnx->engine）

​ 以数字识别模型Densenet为例。

​ 首先，Keras模型需要包含图结构以及模型参数，利用keras2onnx工具，将keras模型转存为onnx模型；

import keras
import keras2onnx
import onnx
from tensorflow.keras.models import load_model

model = load_model('./models/densenet_num18.h5')
onnx_model = keras2onnx.convert_keras(model, model.name)
temp_model_file = './models/densenet_num18.onnx'
onnx.save_model(onnx_model, temp_model_file)

​ 然后有两种方法将onnx模型转换为TensorRT模型；

​ 法一：利用TensorRT自带工具trtexec将onnx模型转化为engine模型。

cd /bin

# 利用trtexec工具转换onnx模型为TensorRT可加载的engine模型
./trtexec \
--onnx=/densenet_num18.onnx \ 
--shapes=the_input:1x32x148x1 \
--workspace=4096 \
--saveEngine=/densenet_num18.engine

最后加载engine模型并推理加速。

import os
import time
import numpy as np
import tensorrt as trt
from PIL import Image
import pycuda.driver as cuda
import pycuda.autoinit

class load_engine_inference(object):
    def __init__(self, file_path):
        self.engine = self.loadEngine2TensorRT(file_path)

    def loadEngine2TensorRT(self, filepath):
        G_LOGGER = trt.Logger(trt.Logger.WARNING)
        # 反序列化引擎
        with open(filepath, "rb") as f, trt.Runtime(G_LOGGER) as runtime:
            engine = runtime.deserialize_cuda_engine(f.read())
            return engine

    def do_inference(self, img_path, do_print=False):
        img = Image.open(img_path)
        img = img.convert('L')
        X = img.reshape([1, 32, width, 1])

        input = X

        # start_time = time.time()
        output = np.empty((1, 18, 19), dtype=np.float32)

        #创建上下文
        self.context = self.engine.create_execution_context()
        # 分配内存
        d_input = cuda.mem_alloc(1 * input.size * input.dtype.itemsize)
        d_output = cuda.mem_alloc(1 * output.size * output.dtype.itemsize)
        bindings = [int(d_input), int(d_output)]

        # pycuda操作缓冲区
        self.stream = cuda.Stream()

        # 将输入数据放入device
        start_time = time.time()
        self.pred_img(input, d_input, bindings, output, d_output)
        end_time = time.time()
        # 线程同步
        self.stream.synchronize()

        self.context.__del__()
        return output, end_time - start_time

    def pred_img(self, input, d_input, bindings, output, d_output):
        cuda.memcpy_htod_async(d_input, input, self.stream)
        # 执行模型
        self.context.execute_async(1, bindings, self.stream.handle, None)
        # 将预测结果从从缓冲区取出
        cuda.memcpy_dtoh_async(output, d_output, self.stream)
        self.stream.synchronize()


img_path = '../tf/num_1_true.bmp'

engine_infer = load_engine_inference('./densenet_num18.engine')

## 第一次推理
res_1, use_time = engine_infer.do_inference(img_path)
print('first inference time: ', np.round((use_time)*1000, 2), 'ms')

法二：直接在代码中通过载入onnx模型并创建engine即可。

# -*- coding:UTF-8 -*-
# @Time    : 2021/5/24
# @Author  : favor
# @Func    : 利用TensorRT对ONNX模型进行加速的推理代码

import os
import sys
import time
import math
import copy
import numpy as np
import tensorrt as trt
import pycuda.autoinit
import pycuda.driver as cuda
from PIL import Image

### 正常TensorRT定义变量
TRT_LOGGER = trt.Logger(trt.Logger.INFO)
a = (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
device = 'cuda:0'

### 分配内存不超过30G
def GiB(val):
    return val * 1 << 30

### 读取onnx模型并构建engine
def build_engine(onnx_path, using_half,engine_file,dynamic_input=True):
    trt.init_libnvinfer_plugins(None, '')
    with trt.Builder(TRT_LOGGER) as builder, builder.create_network(EXPLICIT_BATCH) as network, trt.OnnxParser(network, TRT_LOGGER) as parser:
        builder.max_batch_size = 1 # always 1 for explicit batch
        config = builder.create_builder_config()
        config.max_workspace_size = GiB(4)  # 设置4G的创建engine的显存占用
        if using_half:
            config.set_flag(trt.BuilderFlag.FP16) # 半精度FP6
        # Load the Onnx model and parse it in order to populate the TensorRT network.
        with open(onnx_path, 'rb') as model:
            if not parser.parse(model.read()):
                print ('ERROR: Failed to parse the ONNX file.')
                for error in range(parser.num_errors):
                    print (parser.get_error(error))
                return None
        ### 设置动态输入尺寸设置三个尺寸，依次最小尺寸，最佳尺寸，最大尺寸。
        ### （Batch-size, channel, height, width）输入最好按照该顺序，本人尝试通道数放最后未走通，动态输入会报错。
        if dynamic_input:
            profile = builder.create_optimization_profile();
            profile.set_shape("the_input", (1,1,32,80), (1,1,32,148), (1,1,32,250))
            config.add_optimization_profile(profile)

        return builder.build_engine(network, config)

### 为输入，输出分配内存
def allocate_buffers(engine, is_explicit_batch=False, input_shape=None, output_shape=18):
    inputs = []
    outputs = []
    bindings = []

    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__()

    for binding in engine:
        dims = engine.get_binding_shape(binding)
        ### 此处是动态输入和动态输出所需设置的输入和输出尺寸大小。
        if dims[-1] == -1 and len(dims) == 4:
            assert (input_shape is not None)
            dims[-1] = input_shape
        elif dims[-2] == -1 and len(dims) == 3:
            assert (output_shape is not None)
            dims[-2] = output_shape
        size = trt.volume(dims) * engine.max_batch_size  # 设置推理所需的最大batch-size.
        dtype = trt.nptype(engine.get_binding_dtype(binding))
        # 分配内存
        host_mem = cuda.pagelocked_empty(size, dtype)
        device_mem = cuda.mem_alloc(host_mem.nbytes)
        bindings.append(int(device_mem))
        if engine.binding_is_input(binding):
            inputs.append(HostDeviceMem(host_mem, device_mem))
        else:
            outputs.append(HostDeviceMem(host_mem, device_mem))
    return inputs, outputs, bindings

### 输入图片的预处理，输出预处理过后图片以及该图片预处理过后的宽。
def preprocess_image(imagepath):
    img = Image.open(imagepath)
    img = img.convert('L')
    width, height = img.size[0], img.size[1]
    scale = height * 1.0 / 32
    new_width = int(width / scale)

    img = img.resize([new_width, 32], Image.ANTIALIAS)
    img = np.array(img).astype(np.float32) / 255.0 - 0.5
    X = img.reshape([1, 1, 32, new_width])

    return X, new_width

### 解码预测结果程序，从而得到识别结果以及置信度。
def decode(pred):
    characters = ''.join(
        ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9', ',', '.', '-', '%', '+', '/', '(', ')']) + ' '
    nclass = len(characters)
    char_list = []
    pred_text = pred.argmax(axis=2)[0]
    prob_text = pred.max(axis=2)[0]
    prob_text_second_array = np.transpose(np.sort(pred, axis=2))
    prob_text_second = np.transpose(prob_text_second_array[-2])[0]
    prob_second = 1 - prob_text_second / prob_text

    # char_list 对应的softmax概率列表
    prob_list = []
    prob_sec_list = []

    for i in range(len(pred_text)):
        if pred_text[i] != nclass - 1 and (
                (not (i > 0 and pred_text[i] == pred_text[i - 1]))):
            char_list.append(characters[pred_text[i]])
            prob_list.append(prob_text[i])
            prob_sec_list.append(prob_text_second[i])
    chars = u''.join(char_list)

    # 最小置信度(排除空格)，及对应数字
    try:
        min_p = min(prob_text[np.where(pred_text != nclass - 1)])
        # print(prob_list)
        # print(prob_sec_list)
        min_p_5_6 = 1
        for chi in range(len(char_list)):
            if (char_list[chi] == '5' or char_list[chi] == '6') and prob_list[chi] < min_p_5_6:
                min_p_5_6 = prob_list[chi]
        min_num = 1
        for chi in range(len(char_list)):
            if char_list[chi] in '0123456789' and prob_list[chi] < min_num:
                min_num = prob_list[chi]
    except:
        min_p = 0.0
        min_p_5_6 = 1
        min_num = 1
    return chars, prob_text[np.where(pred_text != nclass - 1)].mean(), min_p, min_num

### 根据输入图片的宽度计算最终模型输出的尺寸大小。（因本人模型为动态输入，动态输出，故需计算输出尺寸大小，并分配输出的内存占用）。
def compute_out_shape(input_shape):
    x2 = input_shape
    x2_ft = math.floor((x2 - 5 + 2 * 2) / 2) + 1
    for i in range(2):
        x2_ft = int(x2_ft/2) if x2_ft % 2 == 0 else int((x2_ft-1)/2)
    return x2_ft

### 根据engine推理代码
def profile_trt(engine, imagepath, batch_size):
    assert (engine is not None)

    ### 确定模型输出尺寸，从而为输入输出分配内存
    input_image, input_shape = preprocess_image(imagepath)
    output_shape = compute_out_shape(input_shape)

    segment_inputs, segment_outputs, segment_bindings = allocate_buffers(engine, True, input_shape, output_shape)

    stream = cuda.Stream()
    with engine.create_execution_context() as context:
        context.active_optimization_profile = 0
        origin_inputshape = context.get_binding_shape(0)
        # 本人此处输入图片的宽度为动态的，故最后一位为动态的，并根据输入图片尺寸进行固定
        if (origin_inputshape[-1] == -1 and len(origin_inputshape) == 4):
            origin_inputshape[-1] = input_shape
            context.set_binding_shape(0, (origin_inputshape))
        # 本人此处模型输出结果为动态 的，为倒数第二位为动态的，可以根据输入图片的宽度确定
        elif (origin_inputshape[-2] == -1 and len(origin_inputshape) == 3):
            origin_inputshape[-2] = output_shape
            context.set_binding_shape(0, (origin_inputshape))

        segment_inputs[0].host = input_image
        start_time = time.time()
        [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in segment_inputs]
        context.execute_async(bindings=segment_bindings, stream_handle=stream.handle)
        [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in segment_outputs]
        stream.synchronize()
        use_time = time.time() - start_time
        infer_out = [out.host for out in segment_outputs]
        results = decode(infer_out[0].reshape([1, output_shape, 19]))

        return results, use_time

if __name__ == '__main__':
    onnx_path = './onnx_models/densenet_num18_rmnode1.onnx'
    usinghalf = True
    batch_size = 1
    imagepath = '../tf/data/num_5_true.bmp'
    engine_file = 'densenet_num18_t_dynamic.engine'
    init_engine = True
    load_engine = True

    ### 初始化并创建engine，根据onnx模型创建engine，该步骤较为费时，故正常会将engine保存下来，方便后期推理。
    if init_engine:
        trt_engine = build_engine(onnx_path, usinghalf, engine_file, dynamic_input=True)
        print('engine built successfully!')
        with open(engine_file, "wb") as f:
            f.write(trt_engine.serialize())
        print('save engine successfully')

    ### 利用上方创建的engine进行推理。平时推理时可以将init_engine设为False，因为engine已保存。
    if load_engine:
        trt.init_libnvinfer_plugins(None, '')
        with open(engine_file, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
            trt_engine = runtime.deserialize_cuda_engine(f.read())

        trt_result, use_time = profile_trt(trt_engine, imagepath, batch_size)
        print(trt_result)

        ### 测试查看前10次，前40次，最后50次的平均推理时长。
        final_50_time = 0
        first_10_time = 0
        first_40_time = 0
        infer_nums = 100
        for i in range(infer_nums):
            trt_result, use_time = profile_trt(trt_engine, imagepath, batch_size)
            if 0 <= i < 10:
                first_10_time += use_time
            elif 10 <= i < 50:
                first_40_time += use_time
            elif 50 <= i <= 100:
                final_50_time += use_time

        print('TensorRT average inference time: first 10: {}ms, medium 40: {}ms,final 50: {}ms'.format(
            np.round(first_10_time * 1000 / 10, 2), np.round(first_40_time * 1000 / (40), 2),
            np.round(final_50_time * 1000 / (50), 2)))

3、利用TensorRT加速Pytorch模型

法一：主体思想：Pytorch->jit->TensorRT（利用Forward框架加速Pytorch模型）

以图片分类模型Resnet18为例。

首先，利用torch自带jit工具，将Pytorch模型转存为jit模型；

import torch
resnet_model_path = '/resnet_state.pth'
model = ResNet()
model.load_state_dict(torch.load(resnet_model_path))
model.eval()
model.cpu()

def TracedModelFactory(file_name, traced_model):
    traced_model.save(file_name)
    traced_model = torch.jit.load(file_name)
    print("filename : ", file_name)
    print(traced_model.graph)

traced_model = torch.jit.trace(model, (a))
TracedModelFactory('resnet_jit.pth', traced_model)

然后利用编译得到的Forward包利用TensorRT加速jit模型。

import torch
import forward
import numpy as np

# 1. 构建Engine
builder = forward.TorchBuilder()

# 读取图片并预处理
img = cv2.imread('need_rec.png')

dummy_inputs = (img)
infer_mode = 'float32'  #  float32 / float16 / int8_calib / int8

builder.set_mode(infer_mode)
# 注：此次构建engine中的dummy_inputs,dummy_inputs为默认最大输入尺寸
engine = builder.build('./models/resnet_jit.pth', dummy_inputs)

outputs = engine.forward(inputs)  # 如果输入是 cuda tensor, 输出也是 cuda tensor. 如果输入是 cpu tensor,  输出也是 cpu tensor

法二：主体思想：Pytorch->onnx->TensorRT（pth->onnx->engine）

首先，利用Pytorch自带onnx转换工具，将Pytorch模型转换为onnx模型；

import torch

resnet_model_path = './models/resnet_state.pth'
model = ResNet()
model.load_state_dict(torch.load(resnet_model_path))
model.eval()
model.cuda()

## 转换ONNX模型并保存
export_onnx_file = "./models/resnet18.onnx"
x = torch.onnx.export(model,  # 待转换的网络模型和参数
                torch.randn(1, 3, 224, 224, device='cpu'), # 虚拟的输入，用于确定输入尺寸和推理计算图每个节点的尺寸
                export_onnx_file,  # 输出文件的名称
                verbose=False,      # 是否以字符串的形式显示计算图
                input_names=["input"]+ ["params_%d"%i for i in range(120)],  # 输入节点的名称，这里也可以给一个list，list中名称分别对应每一层可学习的参数，便于后续查询
                output_names=["output"], # 输出节点的名称
                opset_version=10,   # onnx 支持采用的operator set, 应该和pytorch版本相关，目前我这里最高支持10
                do_constant_folding=True, # 是否压缩常量
                dynamic_axes={"input":{0: "batch_size", 2: "h"}, "output":{0: "batch_size"},} #设置动态维度，此处指明input节点的第0维度可变，命名为batch_size
                )

然后利用TensorRT自带工具trtexec将onnx模型转化为engine模型。（同keras模型转engine模型一致）

后期会增加更多参考链接。