基于C++ LibTorch的YOLOv5图像分割实战：从模型部署到性能优化

1次阅读

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

在工业场景部署 YOLOv5 图像分割模型时，C++ 开发者常遇到三个典型问题：

ABI 兼容性问题 ：PyTorch Python 训练的模型导出时，可能因编译器版本差异导致 C ++ 加载失败。我们曾遇到 gcc7.4 编译的 LibTorch 无法加载 clang 编译导出模型的情况
动态尺寸输入处理 ：产线中物品尺寸差异大，直接 resize 会丢失细节。传统方案是填充黑边，但会引入无效计算，在 1920×1080 图像上可能浪费 30% 计算量
显存碎片化 ：持续处理不同尺寸图像时，频繁申请释放显存会导致性能下降。实测显示连续处理 100 张图后，推理延迟增长 200%

方案	平均延迟 (ms)	峰值显存 (MB)	部署复杂度
LibTorch	15.2	1200	中
ONNX Runtime	18.7	1500	低
TensorRT	12.1	900	高

LibTorch 在部署复杂度和性能之间取得较好平衡，特别适合需要快速迭代的场景。

// 加载量化后的 TorchScript 模型
torch::jit::Module module;
module = torch::jit::load("yolov5s-seg.pt");
module.eval();
module.to(torch::kCUDA);

// 开启半精度推理
torch::Tensor input_tensor = torch::rand({1,3,640,640}).to(torch::kCUDA);
torch::NoGradGuard no_grad;
auto output = module.forward({input_tensor.to(torch::kHalf)});

void preprocess(const cv::Mat& img, float* data) {
    // 使用 AVX2 指令集加速归一化
    __m256 mean = _mm256_set1_ps(0.5f);
    __m256 std = _mm256_set1_ps(0.5f);
    for (int i = 0; i < img.rows; ++i) {const uchar* ptr = img.ptr(i);
        for (int j = 0; j < img.cols; j += 8) {
            __m256 pixel = _mm256_cvtepi32_ps(
                _mm256_cvtepu8_epi32(_mm_loadu_si128((__m128i*)(ptr + j))));
            _mm256_store_ps(data + i*img.cols + j, 
                _mm256_div_ps(_mm256_sub_ps(pixel, mean), std));
        }
    }
}

class DoubleBuffer {
    std::queue<torch::Tensor> cpu_queue; 
    std::queue<torch::Tensor> gpu_queue;
    std::mutex mtx;

    void enqueue(const cv::Mat& img) {
        torch::Tensor tensor = /* 预处理 */;
        std::lock_guard<std::mutex> lock(mtx);
        cpu_queue.push(tensor);
    }

    torch::Tensor dequeue() {std::lock_guard<std::mutex> lock(mtx);
        if (!gpu_queue.empty()) {auto tensor = gpu_queue.front();
            gpu_queue.pop();
            return tensor;
        }
        return torch::Tensor();}
};

class CudaMemoryPool {
    std::vector<void*> pool_;
    size_t chunk_size_;
public:
    CudaMemoryPool(size_t chunk, int prealloc=10) {
        chunk_size_ = chunk;
        for (int i=0; i<prealloc; ++i) {
            void* ptr;
            cudaMalloc(&ptr, chunk_size_);
            pool_.push_back(ptr);
        }
    }

    void* allocate() {if (pool_.empty()) {
            void* ptr;
            cudaMalloc(&ptr, chunk_size_);
            return ptr;
        }
        auto ptr = pool_.back();
        pool_.pop_back();
        return ptr;
    }
};