手把手教你用UNet训练私有数据集（附PyTorch代码解析）

最新推荐文章于 2025-06-18 15:23:05 发布

原创最新推荐文章于 2025-06-18 15:23:05 发布 · 1.3k 阅读

30 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #人工智能 #python #其他

文章目录

一、数据准备的五大核心步骤（必看！）

1. 数据标注的正确姿势

推荐使用Label Studio（开源神器）
标注文件要保存为PNG格式（注意：必须是单通道！）
类别索引从0开始编号（比如背景0/目标1）
错误示范警告：见过有人用jpg存标注图，结果精度直接掉5%！

2. 数据集结构标准模板

my_dataset/
├── images/          # 原始图像
│   ├── 0001.png
│   └── 0002.png
├── masks/           # 标注图像
│   ├── 0001.png
│   └── 0002.png
└── splits.json      # 数据集划分文件

3. 数据增强的黑科技配方

transform = A.Compose([
    A.RandomRotate90(p=0.5),
    A.HorizontalFlip(p=0.5),
    A.VerticalFlip(p=0.5),
    A.GridDistortion(p=0.2),           # （网格扭曲超好用！）
    A.RandomBrightnessContrast(p=0.3), # 亮度对比度调整
    A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))
])

二、魔改UNet的三大秘诀（效果立竿见影！）

1. 深度监督改造方案

class DeepSupervisionUNet(nn.Module):
    def forward(self, x):
        # 各层解码器输出
        outputs = []
        for decoder in self.decoders:
            x = decoder(x)
            outputs.append(x)
        return outputs  # 返回多尺度预测结果

2. 注意力机制加持

class AttentionGate(nn.Module):
    def __init__(self, F_g, F_l):
        super().__init__()
        self.W_g = nn.Conv2d(F_g, F_l, 1)
        self.psi = nn.Sequential(
            nn.Conv2d(F_l, 1, 1),
            nn.Sigmoid()
        )
        
    def forward(self, g, x):
        # g是深层特征，x是浅层特征
        g_conv = self.W_g(g)
        return x * self.psi(g_conv + x)  # 注意力权重相乘

3. 损失函数调优公式

Dice Loss + Focal Loss组合拳：

class HybridLoss(nn.Module):
    def __init__(self, alpha=0.8):
        super().__init__()
        self.alpha = alpha
        
    def forward(self, pred, target):
        dice_loss = 1 - (2*torch.sum(pred*target) + 1e-6)/(torch.sum(pred)+torch.sum(target)+1e-6)
        focal_loss = -target*(1-pred)**2 * torch.log(pred+1e-6)
        return self.alpha*dice_loss + (1-self.alpha)*focal_loss.mean()

三、训练过程中的六个避坑指南

学习率设置玄学（实测有效！）
- 初始lr=3e-4，每10epoch衰减0.5
- 使用OneCycle策略效果更佳
Batch Size的隐藏陷阱
- 小batch（<=8）建议用GroupNorm代替BatchNorm
- 显存不够时尝试梯度累积

早停机制的正确打开方式

if val_loss < best_loss:
    best_loss = val_loss
    patience = 0
    torch.save(model.state_dict(), 'best_model.pth')
else:
    patience += 1
    if patience > 10:  # 连续10次无改善就停止
        break

类别不平衡的破解之道

使用样本加权采样器
在损失函数中增加类别权重

weights = torch.tensor([0.1, 0.9]).cuda()  # 背景:目标=1:9
criterion = nn.CrossEntropyLoss(weight=weights)

四、效果可视化神器推荐

1. 混淆矩阵热力图

from sklearn.metrics import confusion_matrix
import seaborn as sns

cm = confusion_matrix(y_true, y_pred)
sns.heatmap(cm, annot=True, fmt='d')

2. 特征图可视化技巧

# 获取中间层特征
activation = {}
def get_activation(name):
    def hook(model, input, output):
        activation[name] = output.detach()
    return hook

model.encoder[3].register_forward_hook(get_activation('conv3'))

五、常见问题急救包（遇到问题先看这里！）

Q1：预测结果全是背景怎么办？

检查标注索引是否正确
尝试调整损失函数权重
增加数据增强的强度

Q2：显存爆炸怎么破？

使用更小的输入尺寸
尝试混合精度训练

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

Q3：验证集波动太大？

增加验证集样本量
使用指数移动平均（EMA）

class EMA():
    def __init__(self, model, decay=0.999):
        self.model = model
        self.decay = decay
        self.shadow = {}
        self.params = dict(model.named_parameters())
        
    def update(self):
        for name, param in self.params.items():
            self.shadow[name] = self.decay*self.shadow[name] + (1-self.decay)*param.data

六、完整PyTorch代码模板（可直接运行！）

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from dataset import SegmentationDataset
from model import UNet
from losses import HybridLoss

# 超参数配置
config = {
    'lr': 3e-4,
    'batch_size': 8,
    'epochs': 100,
    'input_size': 256,
    'num_classes': 2
}

# 初始化
model = UNet(config['num_classes']).cuda()
optimizer = optim.AdamW(model.parameters(), lr=config['lr'])
criterion = HybridLoss()
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min')

# 数据加载
train_set = SegmentationDataset('data/train', transform)
train_loader = DataLoader(train_set, batch_size=config['batch_size'], shuffle=True)

# 训练循环
for epoch in range(config['epochs']):
    model.train()
    for images, masks in train_loader:
        images = images.cuda()
        masks = masks.cuda()
        
        outputs = model(images)
        loss = criterion(outputs, masks)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    
    # 验证步骤
    model.eval()
    with torch.no_grad():
        val_loss = 0
        for val_images, val_masks in val_loader:
            # ... 计算验证损失
        scheduler.step(val_loss)