使用TensorRT进行加速推理（示例＋代码）

松小白song

已于 2024-07-05 09:02:08 修改

阅读量5.5k

点赞数 24

CC 4.0 BY-SA版权

文章标签： python 深度学习人工智能

于 2024-07-04 15:58:32 首次发布

本文链接：https://blog.csdn.net/weixin_47151388/article/details/140181530

前言

TensorRT 是 NVIDIA 开发的一款高性能深度学习推理引擎，旨在优化神经网络模型并加速其在 NVIDIA GPU 上的推理性能。它支持多种深度学习框架，并提供一系列优化技术，以实现更高的吞吐量和更低的延迟。

一、TensorRT简介

TensorRT（NVIDIA Tensor Runtime）是由 NVIDIA 开发的一款高性能深度学习推理库，用于在 NVIDIA GPU 上进行高效的深度学习推理。它可以优化深度学习模型并将其部署在生产环境中，以实现低延迟和高吞吐量的推理任务。

1.1TensorRT 的主要特点

模型优化：
- 层融合：将多个层融合为一个层以减少内存访问和计算开销。
- 权重量化：将浮点数权重转换为低精度（如 INT8 或 FP16）以减少模型大小和加快计算速度。
- 内存优化：优化内存使用以减少内存带宽和提高数据传输效率。
高效推理：
- 异构计算：利用 GPU 的高并行计算能力进行高效推理。
- 批处理推理：支持批处理输入，提高 GPU 使用效率。
- 动态输入形状：支持动态输入形状，灵活处理不同大小的输入。
易于集成：
- 支持多种深度学习框架：如 TensorFlow、PyTorch、ONNX 等。
- 多语言支持：提供 C++ 和 Python API，方便开发和集成。
- 插件机制：支持自定义层和操作，通过插件机制扩展 TensorRT 的功能。
灵活性和可扩展性：
- 网络定义 API：允许用户通过 API 手动构建和调整深度学习网络。
- 混合精度推理：支持 FP32、FP16 和 INT8 的混合精度计算，兼顾性能和精度。

1.2TensorRT 的工作流程

导出模型：
- 将训练好的模型从深度学习框架（如 PyTorch、TensorFlow）导出为 ONNX 格式。
解析和构建引擎：
- 使用 TensorRT 解析 ONNX 模型，创建网络定义，并进行优化（如层融合、权重量化）。
- 构建 TensorRT 引擎，这是一个高度优化的二进制文件，可以在 GPU 上高效运行。
推理：
- 加载 TensorRT 引擎并创建执行上下文。
- 为输入和输出分配内存缓冲区。
- 将输入数据复制到 GPU，执行推理，并将输出结果从 GPU 复制回主机。

二、具体示例

这个示例展示了如何使用 TensorRT 将一个预训练的 ResNet-18 模型从 PyTorch 导出为 ONNX 格式，并将其转换为 TensorRT 引擎，最后在 GPU 上进行高效的推理。

2.1代码

import torch
import torchvision.models as models
import onnx
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np

# 1. 导出 ResNet-18 模型为 ONNX 格式
model = models.resnet18(pretrained=True).eval()  # 加载预训练的 ResNet-18 模型，并设置为评估模式
dummy_input = torch.randn(1, 3, 224, 224)  # 创建一个随机输入张量，形状为 (1, 3, 224, 224)
torch.onnx.export(model, dummy_input, "resnet18.onnx", verbose=True)  # 导出模型为 ONNX 格式

# 2. 使用 TensorRT 将 ONNX 模型转换为 TensorRT 引擎
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)  # 创建 TensorRT 日志记录器

def build_engine(onnx_file_path, shape=(1, 3, 224, 224)):
    with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.OnnxParser(network, TRT_LOGGER) as parser:
        config = builder.create_builder_config()
        config.max_workspace_size = 1 << 30  # 设置最大工作空间大小为 1 GB
        builder.max_batch_size = 1  # 设置最大批处理大小为 1

        # 读取 ONNX 模型文件
        with open(onnx_file_path, 'rb') as model:
            if not parser.parse(model.read()):
                print('Failed parsing ONNX file.')
                for error in range(parser.num_errors):
                    print(parser.get_error(error))
                return None

        network.get_input(0).shape = shape  # 设置输入形状
        engine = builder.build_engine(network, config)  # 构建 TensorRT 引擎
        return engine

engine = build_engine("resnet18.onnx")  # 构建 TensorRT 引擎

# 3. 加载 TensorRT 引擎并进行推理
def allocate_buffers(engine):
    inputs = []
    outputs = []
    bindings = []
    stream = cuda.Stream()  # 创建 CUDA 流

    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)  # 分配页锁定内存
        dev_mem = cuda.mem_alloc(host_mem.nbytes)  # 分配设备内存
        bindings.append(int(dev_mem))
        if engine.binding_is_input(binding):
            inputs.append((host_mem, dev_mem))  # 添加输入绑定
        else:
            outputs.append((host_mem, dev_mem))  # 添加输出绑定
    return inputs, outputs, bindings, stream

inputs, outputs, bindings, stream = allocate_buffers(engine)  # 分配缓冲区
context = engine.create_execution_context()  # 创建执行上下文

# 4. 推理函数
def infer(context, bindings, inputs, outputs, stream):
    [cuda.memcpy_htod_async(inp[1], inp[0], stream) for inp in inputs]  # 将输入数据从主机复制到设备
    context.execute_async(bindings=bindings, stream_handle=stream.handle)  # 异步执行推理
    [cuda.memcpy_dtoh_async(out[0], out[1], stream) for out in outputs]  # 将输出数据从设备复制到主机
    stream.synchronize()  #