30分钟上手U-2-Net：从标注到推理的自定义数据集训练全流程

你是否曾因开源模型无法精准识别特定物体而困扰？想让AI自动勾勒产品轮廓却受限于通用数据集？本文将带你用U-2-Net实现专属分割模型，从数据标注到模型部署全程可视化操作，无需深厚机器学习背景也能快速上手。完成后你将获得：自定义物体的像素级分割能力、可复用的训练流程模板、模型优化核心参数调优指南。

项目基础与环境准备

U-2-Net（U Square Net）是由秦学斌等人提出的显著目标检测（Salient Object Detection, SOD）模型，通过嵌套U型结构实现高精度的前景提取。项目核心文件结构如下：

• 模型架构定义：model/u2net.py
• 训练入口脚本：u2net_train.py
• 数据加载模块：data_loader.py
• 官方使用说明：README.md

环境配置步骤

1. 克隆项目仓库

git clone https://gitcode.com/gh_mirrors/u2/U-2-Net
cd U-2-Net

2. 安装依赖库（完整列表见requirements.txt）

pip install torch==0.4.0 torchvision==0.2.1 opencv-python scikit-image numpy==1.15.2

3. 下载预训练模型（可选，用于迁移学习）

# 执行模型权重下载脚本
python setup_model_weights.py

数据集构建：从原始图像到标注文件

数据采集规范

优质数据集是模型效果的基础，建议遵循以下标准：

• 图像分辨率不低于320×320像素
• 目标占比不小于图像面积的30%
• 背景多样性：包含目标可能出现的各种场景
• 样本数量：建议至少200张图像（可参考test_data/test_images/中的示例）

标注工具选择与操作

推荐使用LabelMe进行标注（需额外安装pip install labelme），标注流程：

1. 启动标注工具：labelme --flags labels.txt
2. 用多边形工具勾勒目标轮廓
3. 保存为JSON格式（包含图像路径和掩码信息）

标注完成后，通过以下脚本转换为U-2-Net兼容格式：

# 转换LabelMe JSON为图像掩码
python labelme2mask.py --input_dir ./raw_annotations --output_dir ./train_data/gt_aug

数据集目录结构需符合数据加载器要求：

train_data/
├── DUTS/
│   └── DUTS-TR/
│       ├── im_aug/       # 训练图像（JPG格式）
│       └── gt_aug/       # 标注掩码（PNG格式）

训练配置与参数优化

核心参数设置

修改训练脚本中的关键参数：

# 第51-53行：数据集路径
data_dir = os.path.join(os.getcwd(), 'train_data' + os.sep)
tra_image_dir = os.path.join('DUTS', 'DUTS-TR', 'DUTS-TR', 'im_aug' + os.sep)
tra_label_dir = os.path.join('DUTS', 'DUTS-TR', 'DUTS-TR', 'gt_aug' + os.sep)

# 第60-62行：训练超参数
epoch_num = 100          # 建议从50开始，根据验证集调整
batch_size_train = 4     # 根据GPU显存调整（12G显存可设为12）
save_frq = 500           # 每500迭代保存一次模型

数据增强策略

U-2-Net提供多种数据增强方式（定义于数据加载器）：

• RescaleT(320)：等比例缩放至320×320
• RandomCrop(288)：随机裁剪288×288区域
• ToTensorLab(flag=0)：图像归一化与张量转换

增强组合示例：

transforms.Compose([
    RescaleT(320),
    RandomCrop(288),
    ToTensorLab(flag=0)
])

迁移学习配置

使用预训练权重加速收敛：

# 加载基础模型权重
net.load_state_dict(torch.load('saved_models/u2net/u2net.pth'), strict=False)
# 冻结特征提取层
for param in list(net.parameters())[:-20]:
    param.requires_grad = False

模型训练与监控

启动训练进程

执行训练命令，建议使用GPU加速：

# 单GPU训练
python u2net_train.py

# 多GPU分布式训练
python -m torch.distributed.launch --nproc_per_node=2 u2net_train.py

训练过程中会输出损失值变化：

l0: 0.682345, l1: 0.521876, l2: 0.410239, l3: 0.356782, l4: 0.310293, l5: 0.287645, l6: 0.265432
[epoch:  12/100, batch:   120/ 2000, ite: 1200] train loss: 0.325, tar: 0.287

训练可视化

通过TensorBoard监控训练过程：

tensorboard --logdir=./runs --port=6006

重点关注：

• 损失曲线：训练/验证损失应持续下降并趋于稳定
• 预测样本：每10个epoch可视化验证集结果
• 学习率：使用余弦退火策略时观察LR变化

推理部署与效果评估

单张图像推理

使用测试脚本进行推理：

# 推理单张图像
python u2net_test.py --model_path saved_models/u2net/epoch_50.pth \
                     --input_image test_data/test_images/girl.png \
                     --output_dir test_data/u2net_results/

推理结果将保存为PNG格式的掩码图像，可通过以下代码合成前景：

# 合成透明背景图像
python composite.py --image_path test_data/test_images/girl.png \
                    --mask_path test_data/u2net_results/girl.png \
                    --output_path composite_result.png

批量处理与性能优化

针对大规模数据处理，可修改推理脚本中的批量参数：

# 第45行：设置批量大小
batch_size_val = 8

量化模型减小体积并加速推理：

# 模型量化为FP16
torch.save(net.half().state_dict(), 'saved_models/u2net/quantized_model.pth')

常见问题与解决方案

训练过程问题

1. 损失不下降

   • 检查数据标注质量，确保掩码与图像对齐
   • 降低学习率（初始值设为1e-4）
   • 增加数据增强多样性（添加随机旋转）

2. 过拟合现象

   • 增加训练数据量或使用数据增强
   • 添加早停机制（监测验证损失）
   • 调整正则化参数（权重衰减设为1e-5）

3. 内存溢出

   • 减小批量大小（从12降至4）
   • 使用梯度累积（accumulation_steps=4）
   • 启用混合精度训练（torch.cuda.amp）

推理效果优化

1. 边界模糊

   • 后处理使用高斯滤波（sigma=1.5）
   • 调整阈值（置信度设为0.65）
   • 增加输入图像分辨率

2. 小目标漏检

   • 在训练集中增加小目标样本
   • 调整损失函数权重（增加小目标区域权重）
   • 使用多尺度推理（输入图像缩放至不同尺寸）

项目应用与扩展方向

U-2-Net已在多个领域成功应用：

• 背景移除：电商产品图自动抠图
• 人像美化：智能美颜中的面部特征提取
• 工业质检：零件缺陷自动检测

进阶方向：

1. 模型压缩：使用知识蒸馏训练轻量级模型U2NETP
2. 实时推理：转换为ONNX格式部署到边缘设备
3. 半监督学习：结合少量标注数据与大量未标注数据

总结与资源链接

本文详细介绍了U-2-Net自定义数据集训练的全流程，从数据准备到模型部署关键步骤。核心资源：

• 训练脚本：u2net_train.py
• 数据加载：data_loader.py
• 模型定义：model/u2net.py
• 官方文档：README.md

建议收藏本文并关注项目更新，下一篇将介绍如何将训练好的模型部署为Web服务。如有问题，欢迎在项目Issues中交流。