从零构建PyTorch版DeepLabV3+语义分割实战:自定义数据集全流程解析
1. 环境准备与数据标注
在开始构建PyTorch版DeepLabV3+模型之前,我们需要先搭建开发环境。我推荐使用Python 3.8+和PyTorch 2.2+的组合,这个版本组合经过实测稳定性最好。如果你使用Windows系统,建议通过Anaconda创建虚拟环境:
conda create -n deeplab python=3.8
conda activate deeplab
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
数据标注是语义分割项目中最耗时的环节。我尝试过多种标注工具,发现Roboflow的在线标注工具效率最高,特别适合非专业标注人员。以农业果实分割为例,标注时需要注意:
- 果实边缘要精确到像素级
- 不同类别(如成熟果实/未成熟果实)要用不同标签
- 遮挡部分需要根据可见轮廓合理推断
标注完成后,通常会得到YOLO格式的标注文件。我们需要将其转换为VOC格式的语义分割标注。这个转换过程有几个关键点:
def yolo_to_voc(yolo_bbox, img_width, img_height):
"""
将YOLO格式(归一化坐标)转换为VOC格式(像素坐标)
:param yolo_bbox: [class_id, x_center, y_center, width, height]
:return: [x_min, y_min, x_max, y_max]
"""
x_center = float(yolo_bbox[1]) * img_width
y_center = float(yolo_bbox[2]) * img_height
width = float(yolo_bbox[3]) * img_width
height = float(yolo_bbox[4]) * img_height
x_min = int(x_center - width/2)
y_min = int(y_center - height/2)
x_max = int(x_center + width/2)
y_max = int(y_center + height/2)
return [x_min, y_min, x_max, y_max]
2. 数据集构建与预处理
构建符合DeepLabV3+要求的数据集目录结构是项目成功的关键。标准的VOC格式目录应包含:
Seg552/
├── JPEGImages/ # 存放原始图像
├── SegmentationClass/ # 存放语义分割标注图
├── ImageSets/
│ └── Segmentation/ # 存放train.txt/val.txt
└── labels.yaml # 类别定义文件
数据集划分时,我建议采用8:1:1的比例分配训练集、验证集和测试集。特别要注意的是:
- 确保图像和标注文件一一对应
- 验证集应包含所有类别的样本
- 测试集最好来自不同时间或地点的数据
对于图像预处理,DeepLabV3+需要特定的transform管道。这是我经过多次实验验证的有效配置:
from torchvision import transforms
train_transform = transforms.Compose([
tr.RandomScale((0.5, 2.0)), # 多尺度训练
tr.RandomCrop(513), # 随机裁剪
tr.RandomHorizontalFlip(), # 水平翻转
tr.ColorJitter( # 颜色扰动
brightness=0.5,
contrast=0.5,
saturation=0.5
),
tr.Normalize( # 标准化
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]
),
tr.ToTensor() # 转为张量
])
3. 模型配置与修改
DeepLabV3+的PyTorch实现通常需要针对自定义数据集进行多处修改。核心修改点包括:
- 数据集路径配置:在
mypath.py中添加自定义数据集 - 数据加载器修改:创建新的数据集类继承自
torch.utils.data.Dataset - 类别数调整:修改模型输出通道数
这里给出关键的数据集类实现:
class Seg552Dataset(Dataset):
CLASSES = ['background', 'ripe_fruit', 'unripe_fruit']
def __init__(self, root, split='train', transform=None):
self.root = root
self.split = split
self.transform = transform
self.images = self._load_images()
def _load_images(self):
with open(os.path.join(self.root, f'ImageSets/Segmentation/{self.split}.txt')) as f:
return [line.strip() for line in f.readlines()]
def __getitem__(self, idx):
img_name = self.images[idx]
img_path = os.path.join(self.root, 'JPEGImages', img_name + '.jpg')
lbl_path = os.path.join(self.root, 'SegmentationClass', img_name + '.png')
img = Image.open(img_path).convert('RGB')
lbl = Image.open(lbl_path)
if self.transform:
img, lbl = self.transform(img, lbl)
return img, lbl
def __len__(self):
return len(self.images)
特别要注意的是,需要在dataloaders/__init__.py中注册新的数据集,并修改utils.py中的颜色映射函数以适配自定义类别。
4. 模型训练与调优
训练DeepLabV3+时,backbone的选择对结果影响很大。根据我的经验:
- MobileNetV2:训练快、内存占用小,适合移动端部署
- ResNet50/101:精度高但计算量大
- Xception:平衡了速度和精度
这是我常用的训练参数配置:
python train.py \
--backbone mobilenet \
--lr 0.007 \
--workers 4 \
--epochs 100 \
--batch-size 16 \
--gpu-ids 0 \
--checkname deeplab-mobilenet \
--dataset Seg552 \
--use-sbd False \
--loss-type focal \
--weight-decay 1e-4
训练过程中有几个关键技巧:
- 使用学习率warmup:前5个epoch线性增加学习率
- 采用混合精度训练:减少显存占用
- 实施早停机制:验证集指标不再提升时停止训练
对于类别不平衡问题,我推荐使用带权重的交叉熵损失:
class_weight = torch.tensor([0.1, 1.0, 0.8]) # 背景、成熟果实、未成熟果实
criterion = nn.CrossEntropyLoss(weight=class_weight.cuda())
5. 模型测试与部署
模型测试阶段需要特别注意指标的选择。除了常见的mIoU,对于农业应用还应关注:
- 各类别的召回率
- 小目标的检测效果
- 边缘分割的精确度
这是我使用的测试脚本核心部分:
def evaluate(model, dataloader, num_classes):
model.eval()
conf_matrix = np.zeros((num_classes, num_classes))
with torch.no_grad():
for inputs, labels in tqdm(dataloader):
inputs = inputs.cuda()
labels = labels.cuda()
outputs = model(inputs)
preds = outputs.max(1)[1]
# 更新混淆矩阵
mask = (labels >= 0) & (labels < num_classes)
labels = labels[mask]
preds = preds[mask]
conf_matrix += confusion_matrix(
labels.cpu().numpy(),
preds.cpu().numpy(),
labels=np.arange(num_classes)
)
# 计算各类IoU
ious = []
for i in range(num_classes):
tp = conf_matrix[i,i]
fp = conf_matrix[:,i].sum() - tp
fn = conf_matrix[i,:].sum() - tp
iou = tp / (tp + fp + fn + 1e-10)
ious.append(iou)
return np.mean(ious), ious
对于模型部署,我建议:
- 使用TorchScript导出模型
- 实施测试时增强(TTA)提升精度
- 添加后处理(如CRF)优化边缘
实际部署时,这个简单的预处理管道可以显著提升推理速度:
transform = transforms.Compose([
transforms.Resize(512), # 保持长宽比缩放
transforms.CenterCrop(513),
transforms.ToTensor(),
transforms.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]
)
])
6. 常见问题与解决方案
在实际项目中,我遇到过几个典型问题及解决方法:
问题1:标注不一致导致训练震荡
- 现象:验证集指标波动大
- 解决:统一标注标准,对模糊区域明确标注规则
问题2:小目标识别效果差
- 现象:小果实漏检率高
- 解决:
- 使用更高分辨率输入
- 添加针对小目标的损失项
- 采用注意力机制
问题3:模型过拟合
- 现象:训练集指标高但验证集差
- 解决:
- 增加数据增强(如随机旋转、颜色抖动)
- 添加Dropout层
- 使用标签平滑技术
对于计算资源有限的情况,可以尝试这些优化:
# 梯度累积
accum_steps = 4
optimizer.zero_grad()
for i, (inputs, labels) in enumerate(train_loader):
outputs = model(inputs)
loss = criterion(outputs, labels)
loss = loss / accum_steps
loss.backward()
if (i+1) % accum_steps == 0:
optimizer.step()
optimizer.zero_grad()
7. 进阶技巧与优化
经过多个项目的实践,我总结出几个提升模型效果的进阶技巧:
- 知识蒸馏:用大模型指导小模型训练
- 模型集成:融合多个模型的预测结果
- 测试时增强:对输入图像进行多种变换后综合结果
这个TTA实现可以提升约2%的mIoU:
def tta_predict(model, image, scales=[0.5, 0.75, 1.0, 1.25, 1.5]):
preds = []
for scale in scales:
h, w = image.size
new_h, new_w = int(h*scale), int(w*scale)
scaled_img = F.interpolate(image, size=(new_h, new_w), mode='bilinear')
# 水平翻转增强
flipped_img = torch.flip(scaled_img, dims=[3])
with torch.no_grad():
out1 = model(scaled_img)
out2 = model(flipped_img)
out2 = torch.flip(out2, dims=[3])
# 恢复原尺寸
out1 = F.interpolate(out1, size=(h,w), mode='bilinear')
out2 = F.interpolate(out2, size=(h,w), mode='bilinear')
preds.append((out1 + out2)/2)
return torch.mean(torch.stack(preds), dim=0)
对于实际应用场景,还需要考虑:
- 模型量化减小体积
- 使用TensorRT加速
- 开发可视化工具监控预测结果
这是我常用的结果可视化代码,可以直观检查分割效果:
def visualize_prediction(image, pred, gt=None):
plt.figure(figsize=(15,5))
# 原始图像
plt.subplot(1,3,1)
plt.imshow(image)
plt.title('Original Image')
# 预测结果
plt.subplot(1,3,2)
plt.imshow(decode_segmap(pred))
plt.title('Prediction')
# 真实标注(如果有)
if gt is not None:
plt.subplot(1,3,3)
plt.imshow(decode_segmap(gt))
plt.title('Ground Truth')
plt.show()
更多推荐
所有评论(0)