保姆级教程：用PyTorch复现DeeplabV3+语义分割模型（附代码与数据集配置）

weixin_30887919

498人浏览 · 2026-05-26 15:04:44

weixin_30887919 · 2026-05-26 15:04:44 发布

保姆级教程：用PyTorch复现DeeplabV3+语义分割模型（附代码与数据集配置）

在计算机视觉领域，语义分割一直是热门研究方向之一。DeeplabV3+作为Google提出的经典模型，通过改进的encoder-decoder结构和深度可分离卷积，在PASCAL VOC 2012和Cityscapes等基准数据集上取得了领先性能。本文将手把手教你用PyTorch完整复现该模型，从环境搭建到预测可视化，每个步骤都配有可运行的代码片段和详细解释。

1. 环境准备与依赖安装

复现DeeplabV3+需要配置合适的开发环境。推荐使用Anaconda创建独立的Python环境，避免依赖冲突。以下是具体步骤：

conda create -n deeplab python=3.8
conda activate deeplab
pip install torch==1.10.0+cu113 torchvision==0.11.1+cu113 -f https://download.pytorch.org/whl/torch_stable.html
pip install opencv-python pillow matplotlib tqdm

关键依赖说明：

PyTorch 1.10+ ：支持CUDA 11.3的版本
TorchVision ：提供图像预处理工具
OpenCV ：用于图像读写和可视化
Matplotlib ：绘制训练曲线和预测结果

注意：CUDA版本需要与显卡驱动匹配。可通过 nvidia-smi 查看支持的CUDA版本。

硬件建议：

GPU ：至少8GB显存（如RTX 2070及以上）
内存：16GB以上
存储：SSD硬盘加速数据读取

2. 数据集准备与预处理

DeeplabV3+常用PASCAL VOC 2012和Cityscapes数据集。以下以PASCAL VOC为例说明数据处理流程：

2.1 数据集下载与结构

# 数据集目录结构
VOC2012/
├── JPEGImages          # 原始图像
├── SegmentationClass   # 标注图像
├── ImageSets/Segmentation/train.txt  # 训练集列表
└── ImageSets/Segmentation/val.txt    # 验证集列表

2.2 自定义Dataset类

import torch
from torch.utils.data import Dataset
import cv2
import numpy as np

class VOCDataset(Dataset):
    def __init__(self, root, split='train', crop_size=513):
        self.root = root
        self.split = split
        self.crop_size = crop_size
        with open(f"{root}/ImageSets/Segmentation/{split}.txt") as f:
            self.ids = f.read().splitlines()
        
    def __getitem__(self, idx):
        img_path = f"{self.root}/JPEGImages/{self.ids[idx]}.jpg"
        mask_path = f"{self.root}/SegmentationClass/{self.ids[idx]}.png"
        
        img = cv2.imread(img_path)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        mask = cv2.imread(mask_path, 0)
        
        # 随机裁剪和数据增强
        if self.split == 'train':
            h, w = img.shape[:2]
            pad_h = max(self.crop_size - h, 0)
            pad_w = max(self.crop_size - w, 0)
            img = cv2.copyMakeBorder(img, 0, pad_h, 0, pad_w, cv2.BORDER_CONSTANT, value=0)
            mask = cv2.copyMakeBorder(mask, 0, pad_h, 0, pad_w, cv2.BORDER_CONSTANT, value=255)
            
            h, w = img.shape[:2]
            i = np.random.randint(0, h - self.crop_size)
            j = np.random.randint(0, w - self.crop_size)
            img = img[i:i+self.crop_size, j:j+self.crop_size]
            mask = mask[i:i+self.crop_size, j:j+self.crop_size]
            
            if np.random.random() > 0.5:
                img = cv2.flip(img, 1)
                mask = cv2.flip(mask, 1)
        
        # 归一化和转换为Tensor
        img = img.astype(np.float32) / 255.0
        img = torch.from_numpy(img).permute(2, 0, 1)
        mask = torch.from_numpy(mask).long()
        
        return img, mask

3. 模型架构实现

DeeplabV3+的核心创新在于encoder-decoder结构和深度可分离卷积的应用。下面分模块实现：

3.1 深度可分离卷积

class DepthwiseSeparableConv(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1):
        super().__init__()
        self.depthwise = nn.Conv2d(
            in_channels, in_channels, kernel_size, 
            stride, padding, dilation, groups=in_channels
        )
        self.pointwise = nn.Conv2d(in_channels, out_channels, 1)
        
    def forward(self, x):
        x = self.depthwise(x)
        x = self.pointwise(x)
        return x

3.2 ASPP模块

class ASPP(nn.Module):
    def __init__(self, in_channels, out_channels=256):
        super().__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU()
        )
        rates = [6, 12, 18]
        self.convs = nn.ModuleList([
            DepthwiseSeparableConv(in_channels, out_channels, 3, padding=r, dilation=r)
            for r in rates
        ])
        self.pool = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(in_channels, out_channels, 1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU()
        )
        self.project = nn.Sequential(
            nn.Conv2d(5 * out_channels, out_channels, 1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(),
            nn.Dropout(0.5)
        )
        
    def forward(self, x):
        h, w = x.shape[-2:]
        feat1 = self.conv1(x)
        feats = [feat1]
        for conv in self.convs:
            feats.append(conv(x))
        pool = self.pool(x)
        pool = F.interpolate(pool, size=(h, w), mode='bilinear', align_corners=True)
        feats.append(pool)
        out = torch.cat(feats, dim=1)
        return self.project(out)

3.3 完整模型结构

class DeepLabV3Plus(nn.Module):
    def __init__(self, backbone='resnet50', num_classes=21, output_stride=16):
        super().__init__()
        if backbone == 'resnet50':
            self.backbone = resnet50(pretrained=True, replace_stride_with_dilation=[False, False, output_stride==8])
            low_level_channels = 256
        else:
            raise ValueError("Unsupported backbone")
            
        self.aspp = ASPP(2048)
        self.decoder = nn.Sequential(
            nn.Conv2d(low_level_channels, 48, 1),
            nn.BatchNorm2d(48),
            nn.ReLU()
        )
        self.last_conv = nn.Sequential(
            DepthwiseSeparableConv(304, 256, 3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(),
            DepthwiseSeparableConv(256, 256, 3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(),
            nn.Conv2d(256, num_classes, 1)
        )
        
    def forward(self, x):
        h, w = x.shape[-2:]
        # Encoder
        x = self.backbone.conv1(x)
        x = self.backbone.bn1(x)
        x = self.backbone.relu(x)
        x = self.backbone.maxpool(x)
        
        x1 = self.backbone.layer1(x)  # low-level features
        x = self.backbone.layer2(x1)
        x = self.backbone.layer3(x)
        x = self.backbone.layer4(x)   # high-level features
        
        # ASPP
        x = self.aspp(x)
        x = F.interpolate(x, size=(x1.shape[2], x1.shape[3]), mode='bilinear', align_corners=True)
        
        # Decoder
        x1 = self.decoder(x1)
        x = torch.cat([x, x1], dim=1)
        x = self.last_conv(x)
        x = F.interpolate(x, size=(h, w), mode='bilinear', align_corners=True)
        return x

4. 训练策略与超参数调优

4.1 损失函数与优化器

def get_loss_and_optimizer(model, lr=0.007, momentum=0.9, weight_decay=4e-5):
    criterion = nn.CrossEntropyLoss(ignore_index=255)
    optimizer = torch.optim.SGD(
        model.parameters(), 
        lr=lr, 
        momentum=momentum, 
        weight_decay=weight_decay
    )
    return criterion, optimizer

4.2 学习率调度

class PolyLR:
    def __init__(self, optimizer, max_iters, power=0.9):
        self.optimizer = optimizer
        self.max_iters = max_iters
        self.power = power
        self.curr_iter = 0
        
    def step(self):
        self.curr_iter += 1
        lr = self.optimizer.param_groups[0]['lr']
        new_lr = lr * (1 - self.curr_iter / self.max_iters) ** self.power
        for param_group in self.optimizer.param_groups:
            param_group['lr'] = new_lr

4.3 训练循环关键代码

def train_epoch(model, dataloader, criterion, optimizer, lr_scheduler, device):
    model.train()
    total_loss = 0
    for images, masks in tqdm(dataloader):
        images = images.to(device)
        masks = masks.to(device)
        
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, masks)
        loss.backward()
        optimizer.step()
        lr_scheduler.step()
        
        total_loss += loss.item()
    return total_loss / len(dataloader)

5. 模型评估与可视化

5.1 指标计算

def compute_mIoU(preds, labels, num_classes):
    # preds: [B, H, W]
    # labels: [B, H, W]
    ious = []
    for cls in range(num_classes):
        pred_inds = (preds == cls)
        target_inds = (labels == cls)
        intersection = (pred_inds & target_inds).sum().float()
        union = (pred_inds | target_inds).sum().float()
        if union == 0:
            ious.append(float('nan'))
        else:
            ious.append((intersection / union).item())
    return np.nanmean(ious)

5.2 预测可视化

def visualize_prediction(image, pred, gt, save_path=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(gt)
    plt.title("Ground Truth")
    
    plt.subplot(1, 3, 3)
    plt.imshow(pred)
    plt.title("Prediction")
    
    if save_path:
        plt.savefig(save_path)
    plt.close()

6. 常见问题与解决方案

显存不足 ：

减小 batch_size （可降至4或8）

使用混合精度训练：

scaler = torch.cuda.amp.GradScaler()
with torch.cuda.amp.autocast():
    outputs = model(images)
    loss = criterion(outputs, masks)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

训练不收敛 ：
- 检查学习率是否合适（尝试0.001到0.01范围）
- 验证数据预处理是否正确（特别是标签的像素值范围）
- 确保模型输出与标签尺寸匹配
边界分割不准确 ：
- 增加 output_stride 到8（需更多计算资源）
- 在decoder中添加更多卷积层
- 使用更精细的标注数据进行微调