LibtorchSegmentation: 基于Libtorch的C++图像分割库

LibtorchSegmentation简介

LibtorchSegmentation是一个基于Libtorch(PyTorch的C++ API)开发的图像分割库。它为开发者提供了一套高级API,使用户能够轻松地构建、训练和部署各种图像分割模型。该库的主要特点包括:

提供简洁的高级API,只需几行代码即可创建神经网络
支持7种主流的分割模型架构,适用于二分类和多分类任务
包含15种预训练的编码器,可加速模型收敛并提高性能
相比PyTorch CUDA版本,推理速度提升35%以上

这个开源项目旨在为C++开发者提供一个便捷、高效的图像分割解决方案,使其能够快速地将深度学习技术应用到实际项目中。

主要功能与特性

1. 多种分割模型架构

LibtorchSegmentation支持以下7种常用的分割模型架构:

U-Net
FPN (特征金字塔网络)
PAN (路径聚合网络)
PSPNet (金字塔场景解析网络)
LinkNet
DeepLabV3
DeepLabV3+

这些模型涵盖了从经典到最新的各种分割算法,能够满足不同场景下的需求。

2. 丰富的预训练编码器

该库提供了15种预训练的编码器,主要包括:

ResNet系列 (ResNet18, ResNet34, ResNet50等)
ResNeXt系列
VGG系列

所有编码器都预先在ImageNet数据集上进行了训练,可以显著提高模型的收敛速度和性能。用户可以根据具体需求选择合适的编码器。

3. 高效的推理性能

通过优化,LibtorchSegmentation在推理速度上相比PyTorch CUDA版本有显著提升:

CUDA环境下,推理速度提升35%以上
CPU环境下,性能与PyTorch基本持平

这一性能优势使得该库特别适合在实际生产环境中部署。

快速上手

创建分割模型

使用LibtorchSegmentation创建一个分割模型非常简单,只需几行代码:

#include "Segmentor.h"

auto model = UNet(1, // 类别数
                  "resnet34", // 编码器名称
                  "path/to/resnet34.pt" // 预训练权重路径
                  );

这里创建了一个使用ResNet34作为编码器的U-Net模型。用户可以根据需要更换模型架构和编码器。

生成预训练权重

虽然库提供了预训练的编码器,但用户也可以自行生成预训练权重:

import torch
from torchvision import models

model = models.resnet34(pretrained=True)
model.eval()
var = torch.ones((1,3,224,224))
traced_script_module = torch.jit.trace(model, var)
traced_script_module.save("resnet34.pt")

这段Python代码展示了如何导出ResNet34的预训练权重,以供C++代码使用。

训练与预测示例

训练示例

以下代码展示了如何使用LibtorchSegmentation训练一个人像分割模型:

Segmentor<FPN> segmentor;
segmentor.Initialize(0, // GPU ID, -1表示使用CPU
                     512, 512, // 调整图像大小
                     {"background", "person"}, // 类别名称
                     "resnet34",
                     "path/to/resnet34.pt");
segmentor.Train(0.0003, // 初始学习率
                300, // 训练轮数
                4, // 批次大小
                "path/to/dataset",
                ".jpg", // 图像类型
                "path/to/save/segmentor.pt");

这个例子使用FPN模型和ResNet34编码器进行训练,数据集包含背景和人像两个类别。

预测示例

训练完成后,可以使用以下代码进行预测:

cv::Mat image = cv::imread("path/to/image.jpg");
Segmentor<FPN> segmentor;
segmentor.Initialize(0, 512, 512, {"background", "person"},
                     "resnet34", "path/to/resnet34.pt");
segmentor.LoadWeight("segmentor.pt");
segmentor.Predict(image, "person");

这段代码加载训练好的模型,并对输入图像进行人像分割预测。

预测结果示例

自定义数据集训练

LibtorchSegmentation支持用户使用自定义数据集进行训练。主要步骤如下:

使用labelme工具标注图像
将标注后的JSON文件和图像按以下结构组织:

Dataset
├── train
│   ├── xxx.json
│   ├── xxx.jpg
│   └......
├── val
│   ├── xxxx.json
│   ├── xxxx.jpg
│   └......

使用与前面示例类似的代码进行训练,只需将数据集路径替换为自定义数据集的路径即可

此外,库还提供了一些训练技巧,如组合使用Dice损失和交叉熵损失、冻结骨干网络、多步学习率调度等,以进一步提高训练效果。

安装与配置

依赖项

OpenCV 3+
Libtorch 1.7+

安装步骤

配置Libtorch开发环境
安装OpenCV
下载预训练权重文件
构建共享库或静态库:

export Torch_DIR='/path/to/libtorch'
cd build
cmake -DBUILD_SHARED=TRUE ..
make
sudo make install

构建测试:

cd test
mkdir build && cd build
cmake ..
make
./resnet34 ../../voc_person_seg/val/2007_003747.jpg ../../weights/resnet34.pt ../../weights/segmentor.pt