从零开始：使用C++ LibTorch部署YOLOv5图像分割模型的完整指南

1次阅读

没有评论

共计 2933 个字符，预计需要花费 8 分钟才能阅读完成。

YOLOv5 作为目标检测领域的标杆模型，其图像分割变种在工业质检、自动驾驶等场景中表现优异。相比 Python 部署，C++ 结合 LibTorch 能带来显著的性能优势：

执行效率：C++ 原生代码在资源占用和推理速度上普遍快于 Python
生产友好：可直接集成到现有 C ++ 工程，避免跨语言调用的开销
部署灵活：支持无 Python 环境的嵌入式设备部署

推荐使用以下工具链组合：

LibTorch 1.11+（需与训练时 PyTorch 版本严格匹配）
OpenCV 4.5+ 用于图像处理
CMake 3.18+ 作为构建工具

配置示例（Ubuntu 系统）：

wget https://download.pytorch.org/libtorch/cu113/libtorch-cxx11-abi-shared-with-deps-1.11.0%2Bcu113.zip
unzip libtorch*.zip
sudo apt install libopencv-dev

使用官方 export.py 脚本转换 PyTorch 模型为 TorchScript：

下载 YOLOv5 官方代码库
执行转换命令：

export.py --weights yolov5s-seg.pt --include torchscript --img 640

关键参数说明：

--img 必须与训练尺寸一致
建议开启 --optimize 选项优化推理图
输出为 yolov5s-seg.torchscript 文件

基础工程配置示例：

cmake_minimum_required(VERSION 3.18)
project(yolov5_seg)

set(CMAKE_CXX_STANDARD 17)
find_package(OpenCV REQUIRED)

add_executable(demo
  src/main.cpp
)

target_link_libraries(demo
  "${CMAKE_SOURCE_DIR}/libtorch/lib/libtorch.so"
  "${OpenCV_LIBS}"
)

使用 LibTorch 加载序列化模型：

torch::jit::script::Module module;
try {module = torch::jit::load("yolov5s-seg.torchscript");
  module.to(torch::kCUDA);
} catch (const c10::Error& e) {std::cerr << "模型加载失败:" << e.what();
  return -1;
}

符合 YOLOv5 输入规范的预处理流程：

BGR 转 RGB 并归一化
保持长宽比的 resize
填充到正方形尺寸
转换为 CHW 格式

代码实现：

cv::Mat preprocess(cv::Mat img, int target_size) {
  // 转换色彩空间
  cv::cvtColor(img, img, cv::COLOR_BGR2RGB);

  // 计算缩放比例
  float ratio = min(target_size / (img.cols*1.0), target_size / (img.rows*1.0));
  int new_w = img.cols * ratio;
  int new_h = img.rows * ratio;

  // 保持长宽比的 resize
  cv::resize(img, img, cv::Size(new_w, new_h));

  // 填充到正方形
  int top = (target_size - new_h) / 2;
  int bottom = target_size - new_h - top;
  int left = (target_size - new_w) / 2;
  int right = target_size - new_w - left;
  cv::copyMakeBorder(img, img, top, bottom, left, right, 
                    cv::BORDER_CONSTANT, cv::Scalar(114, 114, 114));

  // 转换为 Tensor
  torch::Tensor tensor = torch::from_blob(img.data, {img.rows, img.cols, 3}, torch::kByte);
  tensor = tensor.permute({2, 0, 1}).contiguous();
  tensor = tensor.toType(torch::kFloat).div(255);

  return tensor.unsqueeze(0).to(torch::kCUDA);
}

处理模型输出的关键步骤：

执行模型推理
解析检测框（xywh 格式）
应用 NMS 过滤冗余框
提取并上采样分割 mask

核心代码片段：

auto outputs = module.forward({input_tensor}).toTuple();

torch::Tensor detections = outputs->elements()[0].toTensor();
torch::Tensor masks = outputs->elements()[1].toTensor();

// 处理检测结果
auto conf_mask = detections.select(2, 4).gt(conf_threshold);
detections = detections.index_select(1, conf_mask.nonzero().squeeze());

// 执行 NMS
auto keep = nms(detections.index_select(2, torch::arange(4)), 
              detections.index_select(2, 4), iou_threshold);
detections = detections.index_select(1, keep);

// 处理分割 mask
masks = masks.index_select(2, keep);
masks = torch::sigmoid(masks);

实测优化效果对比（Tesla T4 GPU）：