以yolov5提供一个c++端用libtorch和Torch TensorRT加速推理torchscript模型的案例并注释

注：本文将以 YOLOv5 为例，介绍如何使用 LibTorch 和 Torch TensorRT 对 TorchScript 模型进行加速推理。本文默认读者已经熟悉 YOLOv5 和 TorchScript 的相关知识。

1. 准备工作

在开始之前，需要先安装以下工具：

- PyTorch
- LibTorch
- Torch TensorRT

其中，PyTorch 是用于训练 YOLOv5 模型的框架，而 LibTorch 和 Torch TensorRT 则是用于加速推理的工具。在安装完这些工具之后，需要将训练好的 YOLOv5 模型转换为 TorchScript 格式。

2. 将 YOLOv5 模型转换为 TorchScript 格式

将训练好的 YOLOv5 模型转换为 TorchScript 格式的方法有很多种，这里给出一种比较简单的方法：

import torch
from models.experimental import attempt_load
from utils.general import set_logging
from torch.utils.mobile_optimizer import optimize_for_mobile

def export_torchscript(weights, img_size, device='cpu'):
    set_logging()
    model = attempt_load(weights, map_location=device)
    img = torch.zeros((1, 3, img_size, img_size), device=device)
    model.eval()
    traced_script_module = torch.jit.trace(model, img)
    traced_script_module_optimized = optimize_for_mobile(traced_script_module)
    traced_script_module_optimized.save("yolov5s.torchscript.pt")
    
export_torchscript(weights='yolov5s.pt', img_size=640, device='cpu')

在这个函数中，我们首先加载训练好的 YOLOv5 模型，然后使用 torch.jit.trace 将模型转换为 TorchScript 格式。接着，我们使用 torch.utils.mobile_optimizer.optimize_for_mobile 对模型进行优化，最后将优化后的模型保存到磁盘上。

3. 加载 TorchScript 模型

在 C++ 中加载 TorchScript 模型需要使用 LibTorch，下面是加载模型的代码：

#include <torch/script.h> // One-stop header.

int main(int argc, const char* argv[]) {
  // Load the model.
  torch::jit::script::Module module;
  try {
    // Deserialize the ScriptModule from a file using torch::jit::load().
    module = torch::jit::load("yolov5s.torchscript.pt");
  }
  catch (const c10::Error& e) {
    std::cerr << "error loading the model\n";
    return -1;
  }
  return 0;
}

在这个代码中，我们使用 torch::jit::load 函数加载 TorchScript 模型。如果加载失败，将输出错误信息并返回 -1，否则返回 0。

4. 使用 Torch TensorRT 进行推理

为了加速 TorchScript 模型的推理，我们可以使用 Torch TensorRT。下面是使用 Torch TensorRT 进行推理的代码：

#include <torch/script.h> // One-stop header.
#include <iostream>
#include <memory>
#include <vector>
#include <chrono>
#include <NvInferRuntime.h>

int main(int argc, const char* argv[]) {
  // Load the model.
  torch::jit::script::Module module;
  try {
    // Deserialize the ScriptModule from a file using torch::jit::load().
    module = torch::jit::load("yolov5s.torchscript.pt");
  }
  catch (const c10::Error& e) {
    std::cerr << "error loading the model\n";
    return -1;
  }

  // Create a TensorRT engine from the TorchScript module.
  nvinfer1::IRuntime* runtime = nvinfer1::createInferRuntime(gLogger);
  nvinfer1::ICudaEngine* engine = createCudaEngine(module, runtime, batchSize, kINPUT_BLOB_NAME, kOUTPUT_BLOB_NAME, maxWorkspaceSize);
  if (!engine) {
    std::cerr << "error creating the engine\n";
    return -1;
  }

  // Create a TensorRT execution context.
  nvinfer1::IExecutionContext* context = engine->createExecutionContext();
  if (!context) {
    std::cerr << "error creating the context\n";
    return -1;
  }

  // Prepare inputs and outputs.
  std::vector<float> inputData(batchSize * inputSize * inputSize * 3);
  std::vector<float> outputData(batchSize * outputSize * outputSize * (5 + numClasses));
  void* buffers[] = {inputData.data(), outputData.data()};

  // Run inference.
  auto start = std::chrono::high_resolution_clock::now();
  context->execute(batchSize, buffers);
  auto end = std::chrono::high_resolution_clock::now();
  std::chrono::duration<double, std::milli> elapsed = end - start;
  std::cout << "elapsed time: " << elapsed.count() << " ms\n";

  // Release resources.
  context->destroy();
  engine->destroy();
  runtime->destroy();

  return 0;
}

在这个代码中，我们首先使用 createCudaEngine 函数将 TorchScript 模型转换为 TensorRT engine。接着，我们创建 TensorRT execution context，准备输入和输出数据，并调用 context->execute 进行推理。最后，我们释放资源。

5. 总结

本文介绍了如何使用 LibTorch 和 Torch TensorRT 对 TorchScript 模型进行加速推理。在实际应用中，我们可以根据自己的需求对代码进行修改和优化，以达到更好的性能和效果。