深度学习模型部署TensorRT加速（七）：TensorRT部署图像分类模型

篇章七：TensorRT部署图像分类模型

目录

前言：

一、pytorch构建分类网络

二、tensorrt部署resnet

总结

参考文献

PS:纯粹为学习分享经验，不参与商用价值运作，若有侵权请及时联系！！！

下篇内容预告：

深度学习模型部署TensorRT加速（八）：TensorRT部署目标检测YOLO模型

深度学习模型部署TensorRT加速（九）：TensorRT部署TransFormer模型

---

前言：

使用 TensorRT 和 ONNX 部署图像分类模型并进行推理通常需要以下几个步骤：

准备阶段

1. 环境设置：确保已经安装了 TensorRT、CUDA 和其他必要的依赖。
2. 模型转换：如果模型不是 ONNX 格式，需要将其转换为 ONNX 格式。许多深度学习框架（如 PyTorch、TensorFlow）都提供了工具或方法来进行这一转换。

部署阶段

1. 加载 ONNX 模型：使用 TensorRT 的 API 加载 ONNX /WTS格式的模型。
2. 优化模型：TensorRT 会对模型进行一系列优化，以提高推理速度。这通常包括层融合、精度调整等。
3. 生成引擎：完成优化后，TensorRT 会生成一个用于推理的执行引擎。

推理阶段

1. 数据预处理：将输入图像转换为模型所需的格式和尺寸。
2. 执行推理：将预处理后的数据传递给 TensorRT 引擎，并执行推理。
3. 数据后处理：将推理结果转换为可解释的标签或其他输出。

下列将逐步展示Pytorch代码构建模型部署的具体操作流程。为了更灵活掌握，本章内容展示在Python训练好的权重，如何利用C++进行部署！！！！

一、pytorch构建分类网络

 1） 基于torchvision构建resnet网络，建resnet分类网络，并保存pth权重.

from torchvision.transforms import transforms
import torch
import torchvision.models as models
import struct
transform_train = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
transforms_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

def build_model():
    model = models.resnet18(pretrained=True)
    model = model.eval()
    model = model.cuda()
    torch.save(model, "./resnet18.pth")
if __name__ == '__main__':
    build_model()

2）需要获取中间表示，将权重转化为tensorrt可读的模式(根据情况灵活选择wts或onnx中间表示)

（a) 获得wts文件，获得wts权重格式文件，代码如下：

from torchvision.transforms import transforms
import torch
import torchvision.models as models
import struct
def get_wts(model_path='./resnet18.pth',save_wts_path="./resnet18.wts"):
    net = torch.load(model_path)
    net = net.cuda()
    net = net.eval()
    print('model: ', net)
    # print('state dict: ', net.state_dict().keys())
    tmp = torch.ones(1, 3, 224, 224).cuda()
    print('input: ', tmp)
    out = net(tmp)
    print('output:', out)
    f = open(save_wts_path, 'w')
    f.write("{}\n".format(len(net.state_dict().keys())))
    for k, v in net.state_dict().items():
        print('key: ', k)
        print('value: ', v.shape)
        vr = v.reshape(-1).cpu().numpy()
        f.write("{} {}".format(k, len(vr)))
        for vv in vr:
            f.write(" ")
            f.write(struct.pack(">f", float(vv)).hex())
        f.write("\n")
if __name__ == '__main__':
    get_wts(model_path='./resnet18.pth',save_wts_path="./resnet18.wts")

(b）获得onnx文件（推荐做法）

获得onnx格式文件，代码如下：

from torchvision.transforms import transforms
import torch
import torchvision.models as models
import struct
def get_onnx(model_path='./resnet18.pth',save_onnx_path="./resnet18.onnx"):

    # 定义静态onnx，若推理input_data格式不一致，将导致保存

​    input_data = torch.randn(2, 3, 224, 224).cuda()
​    model = torch.load(model_path).cuda()
​    input_names = ["data"] + ["called_%d" % i for i in range(2)]
​    output_names = ["prob"]
​    torch.onnx.export(
​        model,
​        input_data,
​        save_onnx_path,
​        verbose=True,
​        input_names=input_names,
​        output_names=output_names
​    )

if __name__ == '__main__':
    get_onnx(model_path='./resnet18.pth', save_onnx_path="./resnet18.onnx")

二、tensorrt部署resnet

（a) 基于wts格式采用C++ API 转tensorrt部署 以下使用wts方法，实现引擎engine构建与推理部署，代码如下：

#include <fstream>
#include <iostream>
#include <vector>
#include <NvInfer.h>

using namespace nvinfer1;

// 加载 .wts 文件中的权重到 TensorRT 网络中
// 这里假设有一个名为 loadWeights 的函数，可以根据实际情况自定义一个loadWeights的函数
void loadWeights(const std::string& file, nvinfer1::INetworkDefinition* network) {
    // 实现权重加载逻辑
}

int main() {
    // 创建 TensorRT logger
    nvinfer1::ILogger logger;

    // 创建 TensorRT builder 和 network
    IBuilder* builder = nvinfer1::createInferBuilder(logger);
    INetworkDefinition* network = builder->createNetworkV2(0U);

    // 加载 .wts 模型文件
    loadWeights("ResNet-18.wts", network);

    // 配置 builder 设置
    builder->setMaxBatchSize(1);
    builder->setMaxWorkspaceSize(1 << 20);

    // 创建 TensorRT 引擎
    ICudaEngine* engine = builder->buildCudaEngine(*network);
    if (!engine) {
        std::cerr << "Failed to create TensorRT engine!" << std::endl;
        return -1;
    }

    // 创建执行上下文
    IExecutionContext* context = engine->createExecutionContext();

    // 分配设备内存
    void* deviceInput;
    void* deviceOutput;
    cudaMalloc(&deviceInput, sizeof(float) * 224 * 224 * 3); // 假设输入是 224x224x3
    cudaMalloc(&deviceOutput, sizeof(float) * 1000); // 假设输出是 1000 类

    // 数据预处理（这里假设你有一个名为 preprocess 的函数来处理输入数据）
    // float* inputData = preprocess(inputImage);
    // cudaMemcpy(deviceInput, inputData, sizeof(float) * 224 * 224 * 3, cudaMemcpyHostToDevice);

    // 执行推理
    void* bindings[] = {deviceInput, deviceOutput};
    context->execute(1, bindings);

    // 拷贝输出数据回主机
    float* outputData = new float[1000];
    cudaMemcpy(outputData, deviceOutput, sizeof(float) * 1000, cudaMemcpyDeviceToHost);

    // 后处理（这里假设你有一个名为 postprocess 的函数来处理输出数据）
    // postprocess(outputData);

    // 释放资源
    cudaFree(deviceInput);
    cudaFree(deviceOutput);
    delete[] outputData;
    context->destroy();
    engine->destroy();
    network->destroy();
    builder->destroy();

    return 0;
}

(b) 同样地，基于onnx格式也可以采用C++ API 转tensorrt部署（推荐）

#include <iostream>
#include <fstream>
#include <NvInfer.h>
#include <NvOnnxParser.h>
#include <cuda_runtime_api.h>

using namespace nvinfer1;

int main() {
    // 创建 TensorRT logger
    ILogger logger;

    // 创建 TensorRT builder 和 network
    IBuilder* builder = createInferBuilder(logger);
    INetworkDefinition* network = builder->createNetworkV2(0U);

    // 创建 ONNX 解析器并解析 ONNX 文件
    nvonnxparser::IParser* parser = nvonnxparser::createParser(*network, logger);
    if (!parser->parseFromFile("resnet18.onnx", static_cast<int>(ILogger::Severity::kINFO))) {
        std::cerr << "ERROR: Could not parse the model." << std::endl;
        return -1;
    }

    // 配置 builder 设置
    builder->setMaxBatchSize(1);
    builder->setMaxWorkspaceSize(1 << 20);

    // 创建 TensorRT 引擎
    ICudaEngine* engine = builder->buildCudaEngine(*network);
    if (!engine) {
        std::cerr << "Failed to create TensorRT engine!" << std::endl;
        return -1;
    }

    // 创建执行上下文
    IExecutionContext* context = engine->createExecutionContext();

    // 分配设备内存
    void* deviceInput;
    void* deviceOutput;
    cudaMalloc(&deviceInput, sizeof(float) * 224 * 224 * 3); // 假设输入是 224x224x3
    cudaMalloc(&deviceOutput, sizeof(float) * 1000); // 假设输出是 1000 类

    // 数据预处理（这里假设你有一个名为 preprocess 的函数来处理输入数据）
    // float* inputData = preprocess(inputImage);
    // cudaMemcpy(deviceInput, inputData, sizeof(float) * 224 * 224 * 3, cudaMemcpyHostToDevice);

    // 执行推理
    void* bindings[] = {deviceInput, deviceOutput};
    context->execute(1, bindings);

    // 拷贝输出数据回主机
    float* outputData = new float[1000];
    cudaMemcpy(outputData, deviceOutput, sizeof(float) * 1000, cudaMemcpyDeviceToHost);

    // 后处理（这里假设你有一个名为 postprocess 的函数来处理输出数据）
    // postprocess(outputData);

    // 释放资源
    cudaFree(deviceInput);
    cudaFree(deviceOutput);
    delete[] outputData;
    context->destroy();
    engine->destroy();
    network->destroy();
    builder->destroy();
    parser->destroy();

    return 0;
}

 (c) 事实上为了更好地发挥C++部署的优势，将 ONNX 模型还可以转换为 TensorRT 引擎并保存：

        使用 Python 的 TensorRT 绑定来读取 ONNX 模型并创建一个优化过的 TensorRT 引擎，然后将其保存为 .engine 文件。这样的做法是为了更好地适配引擎推理！！！

• 利用python代码生成.engine文件

  import tensorrt as trt
  
  TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
  
  # 加载 ONNX 模型并创建 TensorRT 引擎
  with trt.Builder(TRT_LOGGER) as builder, builder.create_network(1) as network, trt.OnnxParser(network, TRT_LOGGER) as parser:
      builder.max_workspace_size = 1 << 30  # 1GB
      builder.max_batch_size = 1
  
      # 解析 ONNX 模型
      with open("resnet18.onnx", "rb") as model:
          parser.parse(model.read())
  
      # 创建 TensorRT 引擎
      engine = builder.build_cuda_engine(network)
  
      # 保存引擎到文件
      with open("resnet18.engine", "wb") as f:
          f.write(engine.serialize())
  

• 然后加载，并使用 TensorRT 引擎，利用C++ API进行推理

  #include <fstream>
  #include <iostream>
  #include <vector>
  #include <NvInfer.h>
  #include <cuda_runtime_api.h>
  
  using namespace nvinfer1;
  
  int main() {
      // 创建 TensorRT logger
      ILogger logger;
  
      // 从文件中读取已保存的引擎
      std::ifstream engineFile("resnet18.engine", std::ios::binary);
      std::vector<char> engineData((std::istreambuf_iterator<char>(engineFile)), std::istreambuf_iterator<char>());
      engineFile.close();
  
      // 使用读取的数据创建运行时和引擎
      IRuntime* runtime = createInferRuntime(logger);
      ICudaEngine* engine = runtime->deserializeCudaEngine(engineData.data(), engineData.size(), nullptr);
  
      // 创建执行上下文
      IExecutionContext* context = engine->createExecutionContext();
  
      // ...（与上面的 C++ 示例相同，进行内存分配、数据预处理、推理、数据后处理等）
  
      // 释放资源
      context->destroy();
      engine->destroy();
      runtime->destroy();
  
      return 0;
  }
  

          这样就可以充分发挥开发和部署的各自优势，相互结合。在 Python 环境中轻松地进行模型转换和优化，然后在 C++ 环境中进行高效的推理。

总结：

TensorRT支持Python和C++两种语言进行部署，每种语言都有其独特的优势和适用场景。

Python部署的优势：

1. 快速原型开发： Python是一种高级语言，编写和调试代码相对简单快捷，适用于快速原型开发和迭代。

2. 灵活性： Python拥有丰富的第三方库和框架支持，使得TensorRT与其他Python库（如PyTorch、TensorFlow）无缝集成，可以更方便地进行模型预处理、后处理以及结果可视化等操作。

3. 生态系统： Python拥有庞大的开源社区和资源，可以轻松获取和共享模型、工具和技术。

4. 数据处理和科学计算： Python生态系统中有许多数据处理和科学计算库（如NumPy、Pandas），有助于处理输入数据和解释推理结果。

C++部署的优势：

1. 性能： C++是一种编译型语言，执行效率较高，运行时的性能通常优于Python。特别是在性能敏感的应用中，C++部署可以更好地发挥TensorRT的加速能力。

2. 部署规模： C++程序相对于Python的执行速度更快，对于需要大规模部署的场景，C++部署可以提供更好的性能和效率。

3. 资源消耗： C++部署通常比Python部署占用更少的系统资源，如内存占用更小，适用于嵌入式系统等资源有限的场景。

参考文献：

图像分类实战：mobilenetv2从训练到TensorRT部署（pytorch） - 知乎 (zhihu.com)

(263条消息) tensorRT部署分类网络resnet与性能验证教程(C++)_resnet tensorrt_tangjunjun-owen的博客-CSDN博客

PS:纯粹为学习分享经验，不参与商用价值运作，若有侵权请及时联系！！！

下篇内容预告：

• 深度学习模型部署TensorRT加速（八）：TensorRT部署目标检测YOLO模型

• 深度学习模型部署TensorRT加速（九）：TensorRT部署TransFormer模型