（四）TensorRT | 基于 GPU 端的 Python 推理

最新推荐文章于 2024-04-30 19:24:04 发布

zhangts20

最新推荐文章于 2024-04-30 19:24:04 发布

阅读量2.9k

点赞数 2

分类专栏：模型部署文章标签： python 深度学习机器学习

本文链接：https://blog.csdn.net/Skies_/article/details/126554430

版权

模型部署专栏收录该内容

4 篇文章 1 订阅

订阅专栏

1. TensorRT 的简介和安装

TensorRT 是一种基于英伟达硬件的高性能的深度学习前向推理框架，本文介绍使用 TensorRT 在通用 GPU 上的部署流程。

本地需先安装 CUDA，以 CUDA11.0、TensorRT-8.2.5.1 为例。首先，去官网下载（需先登录）对应的压缩包。Python 安装文件 whl 位于解压后根目录下的 python 文件夹内，pip 安装对应版本即可。

本文主要代码来自 TensorRT 仓库：https://github.com/NVIDIA/TensorRT/tree/main/samples/python/yolov3_onnx。

2. TensorRT 的基本使用

和通用的推理流程类似，本文代码的流程按照：输入图像的预处理，把处理后的速度拷贝到设备（GPU）上，在设备上运行模型推理，将设备上的推理结果拷贝到主机（本地 CPU），针对推理结果的后处理。

2.1 模型转换

本文以 ONNX 为基础，将其转换为基于 TensorRT 推理的 trt 文件格式。关键函数为 runtime.deserialize_cuda_engine，函数原型为：

deserialize_cuda_engine(self: tensorrt.tensorrt.Runtime, serialized_engine: buffer)→ tensorrt.tensorrt.ICudaEngine

可以看到它是类 tensorrt.tensorrt.Runtime 的成员函数，参数只有 serialized_engine 一个。该参数可来自于读取已存在的 trt 文件，或通过序列化 ONNX 模型得到。

2.1.1 读取 trt 文件

with open(engine_file_path, "rb") as f:
    with trt.Runtime(TRT_LOGGER) as runtime:
        return runtime.deserialize_cuda_engine(f.read())

2.1.2 序列化 ONNX 模型

序列化 ONNX 模型的关键函数是 build_serialized_network，函数原型为：

build_serialized_network(self: tensorrt.tensorrt.Builder, network: tensorrt.tensorrt.INetworkDefinition, config: tensorrt.tensorrt.IBuilderConfig)→ tensorrt.tensorrt.IHostMemory

可以看到它是类 tensorrt.tensorrt.Builder 类的成员函数，参数有 network 和 config 两个。函数返回类型和上述 f.read() 的内容类似，作为函数 deserialize_cuda_engine 的参数。

第一个参数 network 的类型为 tensorrt.tensorrt.INetworkDefinition，通过函数 create_network 得到，函数原型为：

create_network(self: tensorrt.tensorrt.Builder, flags: int = 0)→ tensorrt.tensorrt.INetworkDefinition

参数 flags 的内容与 TensorRT 中的显示 Batch 和隐式 Batch 模式有关，相关内容可参考文档。在 Python 中，TensorRT 推荐写法：

network = builder.create_network(
    1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))

第二个参数 config 的类型为 tensorrt.tensorrt.IBuilderConfig，通过函数 create_builder_config 得到，函数原型为：

create_builder_config(self: tensorrt.tensorrt.Builder)→ tensorrt.tensorrt.IBuilderConfig

该函数主要用于设置一些配置项，配置内容可参考文档。

2.1.3 get_engine

把上述两种方式结合起来，得到 get_engine 函数内容，返回反序列化后的模型。

def get_engine(onnx_file_path, engine_file_path=""):
    # 如果不指定 engine_file_path 则通过 build_engine 生成 engine 文件
    def build_engine():
        # 基于 INetworkDefinition 构建 ICudaEngine
        builder = trt.Builder(TRT_LOGGER)
        # 基于 INetworkDefinition 和 IBuilderConfig 构建 engine
        network = builder.create_network(
            1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
        # 构建 builder 的配置对象
        config = builder.create_builder_config()
        # 构建 ONNX 解析器
        parser = trt.OnnxParser(network, TRT_LOGGER)
        # 构建 TensorRT 运行时
        runtime = trt.Runtime(TRT_LOGGER)
        # 参数设置
        config.max_workspace_size = 1 << 28 # 256MiB
        builder.max_batch_size = 1
        # 解析 onnx 模型
        if not os.path.exists(onnx_file_path):
            print(
                f"[INFO] ONNX file {onnx_file_path} not found.")
        print(f"[INFO] Loading ONNX file from {onnx_file_path}.")
        with open(onnx_file_path, "rb") as model:
            print("[INFO] Beginning ONNX file parsing.")
            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 
        # 根据 yolov3.onnx，reshape 输入数据的形状
        network.get_input(0).shape = [1, 3, 608, 608]
        print("[INFO] Completed parsing of ONNX file.")
        print(f"[INFO] Building an engine from {onnx_file_path}.")
        # 序列化模型
        plan = builder.build_serialized_network(network, config)
        # 反序列化
        engine = runtime.deserialize_cuda_engine(plan)
        print("[INFO] Completed creating engine.")
        # 写入文件
        with open(engine_file_path, "wb") as f:
            f.write(plan)
        return engine 

    if os.path.exists(engine_file_path):
        print(f"[INFO] Reading engine from {engine_file_path}.")
        with open(engine_file_path, "rb") as f:
            with trt.Runtime(TRT_LOGGER) as runtime:
                return runtime.deserialize_cuda_engine(f.read())
    else:
        return build_engine()

2.2 输入图像预处理

本文使用 YOLOv3 模型来自 DarkNet，预处理主要包括 resize 和归一化两种。

class PreprocessYOLO:
    def __init__(self, input_resolution):
        self.input_resolution = input_resolution

    def preprocess(self, image_path):
        image_raw, image_resized = self.load_and_resize(image_path)
        image_preprocesed = self.shuffle_and_normalize(image_resized)
        return image_raw, image_preprocesed

    def load_and_resize(self, image_path):
        image_raw = Image.open(image_path)
        new_resolution = (self.input_resolution[1], self.input_resolution[0])
        image_resized = image_raw.resize(new_resolution, resample=Image.CUBIC)
        image_resized = np.array(image_resized, dtype=np.float32, order="C")
        return image_raw, image_resized

    def shuffle_and_normalize(self, image):
        # 归一化
        image /= 255.0
        # (w,h,c) -> (c,h,w)
        image = np.transpose(image, [2, 0, 1])
        # (c,h,w) -> (n,c,h,w)
        image = np.expand_dims(image, axis=0)
        image = np.array(image, dtype=np.float32, order="C")
        return image

2.3 推理

在执行模型推理前，首先要在本地 CPU 和设备 GPU 上分配内存。

def allocate_buffers(engine):
    inputs   = []
    outputs  = []
    bindings = []
    stream = cuda.Stream()
    for binding in engine:
        size  = trt.volume(
            engine.get_binding_shape(binding)) * engine.max_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, stream

然后，在执行推理时首先将图像数据从 CPU 拷贝到 GPU，然后执行推理，最后将推理结果从 GPU 拷贝到 CPU。

def do_inference(context, bindings, inputs, outputs, stream):
    # 将输入数据从主机拷贝到设备
    [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]
    # 推理
    context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)
    # 将输出数据从设备拷贝到主机
    [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]
    # 同步流
    stream.synchronize()
    # 仅返回主机上的输出
    return [out.host for out in outputs]