简介面向语义分割入门与人物抠图应用这份资源提供了基于PyTorch的DeeplabV3完整实战代码与配套数据。项目使用PyTorch自带预训练模型涵盖ResNet50、ResNet101、MobileNetV3_Large三种骨干网络完整演示二分类语义分割的训练、验证与预测流程同时整合Wandb可视化、交叉熵与Dice Loss组合帮助对比不同模型和损失函数的效果。压缩包共3453个文件大小约990.96MB主体为3411张PNG人物图像另有15个Python脚本、20个pyc缓存、Dockerfile、Shell脚本及配置文件可据此快速搭建环境并复现人物抠图。已有1260人学习适合具备一定Python与深度学习基础、希望直接跑通分割项目的开发者后续可通过更换骨干网络或调整Loss进一步优化分割精度。1. 人物抠图为什么选 deeplabv3torchvision 开箱即用的二分割方案做人物抠图最省事的路径其实不是训练一个完整的 U-Net而是直接拿 torchvision 自带的 deeplabv3 改输出通道数。这个 zip 里装的就是一套能直接跑通的工程ResNet50 做骨架、接 ASPP 模块、只预测前景背景两个类别配合交叉熵和 dice_loss 的组合损失训练完用一张图就能验证分割结果。如果你手头有带人物的图片数据集想快速出一个可用的抠图模型又不想从零搭网络结构这套方案比你想的还要直接——torchvision 把模型权重和预处理都封装好了你要做的只是把分类头从 21 类改成 2 类然后把数据喂进去。适合谁已经在用 PyTorch、想落地语义分割但还没碰过 deeplab 系列的人以及那些被网上各种魔改分割模型搞晕、想先跑通一个标准方案再迭代的工程师。2. 数据准备与模型选型从 resnet50 到 mobilenet_v3 的取舍2.1 训练数据怎么组织图片和 mask 的配对方式打开这个 zip 能看到一串以数字命名的 PNG比如 00315.png、00458.png、00916.png还有 output.jpg 和 README.md。这里的常规安排是数字命名的 PNG 是原始图像对应同名的 mask 图放在另一个目录里mask 是黑白二值图白色是人物前景黑色是背景。如果你打开 README 确认了目录结构会发现它并没有把 mask 单独拆出来而是靠命名规则配对这在小型数据集里很常见——省去 JSON 标注文件的维护成本写个 Dataset 类按文件名自动找对应 mask 就行。import os from torch.utils.data import Dataset from PIL import Image class PersonSegDataset(Dataset): def __init__(self, img_dir, mask_dir, img_size(320, 320)): self.img_dir img_dir self.mask_dir mask_dir self.img_size img_size self.names [f for f in os.listdir(img_dir) if f.endswith(.png)] def __len__(self): return len(self.names) def __getitem__(self, idx): name self.names[idx] img Image.open(os.path.join(self.img_dir, name)).convert(RGB) mask Image.open(os.path.join(self.mask_dir, name.replace(.png, _mask.png))).convert(L) img img.resize(self.img_size, Image.BILINEAR) mask mask.resize(self.img_size, Image.NEAREST) return img, mask这段代码里 mask 的 resize 必须用 Image.NEAREST不能用 BILINEAR。原因很简单mask 是离散的类别标签双线性插值会在边缘产生 0.5 这种中间值转 tensor 后就成了 0 和 1 之外的第三类训练时 dice_loss 直接算出个不伦不类的数字。图像本身用 BILINEAR 没问题因为图片是连续的颜色信号插值不会引入语义错误。输入尺寸默认 320x320这个值是 deeplabv3 在 Cityscapes 预训练时的常见尺寸之一。如果你显卡显存够大改成 512 甚至 640 能明显提升边缘质量代价是训练时间几乎翻倍。我的建议是第一轮先用 320 跑通流程确认 loss 在下降再回头调尺寸。2.2 加载预训练模型的两个关键改动torchvision 里 deeplabv3 有 resnet50、resnet101 和 mobilenet_v3_large 三个变体。resnet50 是速度和精度的平衡点resnet101 精度高一些但显存压力大mobilenet_v3_large 适合推理部署。加载时不能只改 num_classes还要处理 aux_classifier——deeplabv3 在训练时有一个辅助分类头如果这个头没改前向传播返回的 dict 里 aux 输出维度还是 21loss 计算会直接报错。import torch from torchvision import models def build_model(num_classes2, pretrainedTrue): model models.segmentation.deeplabv3_resnet50( weightsmodels.DeepLabV3_ResNet50_Weights.COCO_WITH_VOC_LABELS_V1 ) # 主分类头输入通道 2048输出改成 2 类 model.classifier[4] torch.nn.Conv2d(2048, num_classes, kernel_size1) # 辅助分类头deeplabv3 训练时要用不改成 2 类会报维度错 model.aux_classifier[4] torch.nn.Conv2d(256, num_classes, kernel_size1) return model注意这里给的是 resnet50 的通道数换成 mobilenet_v3_large 时主分类头输入是 960aux 是 128。改错通道数的话PyTorch 会在第一次前向时给出 shape mismatch 的报错顺着报错改就行。另外 weights 参数在 torchvision 新版本里推荐用枚举类型而不是字符串老代码里pretrainedTrue的写法在 0.13 之后会告警但不影响运行。2.3 三个变体的选择决策表模型主分类头输入通道显存占用320x320推理速度适用场景deeplabv3_resnet502048约 6-7G中等精度优先训练资源充足deeplabv3_resnet1012048约 9-10G较慢追求极致精度边缘细节复杂deeplabv3_mobilenet_v3_large960约 3-4G快低显存环境后续要转 onnx3. 损失函数组合交叉熵与 dice_loss 的加权策略3.1 为什么单用交叉熵在人物分割上不够人物抠图这个场景有个天然特点大部分图里人物只占画面的 30%-50%背景占比高。交叉熵对每个像素一视同仁模型只要把全部预测成背景loss 就已经很低了。训练初期你会看到 loss 掉得很快但 val 的 IoU 涨不动——这就是类别不平衡在作怪。有人选择给背景类降低权重但权重系数调起来很玄学每个数据集都要重新试。dice_loss 的思路不同它直接优化分割区域的重叠程度对前景小目标更敏感。把两者加起来是分割任务里非常成熟的做法交叉熵保证每个像素的梯度信号稳定dice_loss 把优化方向往「前景区域得分的提升」上拉。def dice_loss(pred, target, smooth1.0): pred torch.sigmoid(pred) # 二分类先过 sigmoid 归一化到 0-1 pred pred.contiguous().view(-1) target target.contiguous().view(-1) intersection (pred * target).sum() return 1 - (2.0 * intersection smooth) / (pred.sum() target.sum() smooth)这段 dice_loss 的写法用的是二分类版本的公式配合交叉熵一起算总 loss。smooth 参数的用意是防止分母为 0数值上相当于给分子分母同时加了一个小常数在 mask 全黑或者全白的情况下也能稳定输出。值得注意的一点是这里的 pred 在使用前先过了 sigmoid而交叉熵部分用的是带 logits 的版本——如果直接把 BCEWithLogitsLoss 的结果和 dice_loss 相加相当于对同样的输出做了两次 sigmoid梯度会乱掉。正确做法是交叉熵用nn.CrossEntropyLoss配合模型原始输出dice_loss 单独用 sigmoid 后的值。实际组合时我一般这样写训练循环里的 loss 计算from torch.nn import CrossEntropyLoss criterion CrossEntropyLoss() # 训练循环内 outputs model(images) # outputs[out] 是主输出 out outputs[out] loss_ce criterion(out, masks.squeeze(1)) # 交叉熵mask 是 (N,H,W) 的长整型 pred_sigmoid torch.softmax(out, dim1)[:, 1, :, :] # 取前景类的概率 loss_dice dice_loss(pred_sigmoid, masks.float()) loss loss_ce 0.5 * loss_dice3.2 权重系数怎么定从 0.5 到 1.0 的调参路径交叉熵和 dice_loss 的配比没有标准答案。dice_loss 的值域在 0 到 1 之间交叉熵在有 2 个类别时通常在 0.2 到 0.7 之间量级差不多所以初始权重可以设成 1:1。如果训练中发现前景区域经常漏检把 dice_loss 的权重提到 0.7 甚至 1.0如果发现训练后期 loss 震荡厉害、收敛不稳说明 dice_loss 梯度过大降回 0.3 左右。这个调整要放在训练早期完成因为后期模型已经定型loss 权重的改变只会在原地抖动。还有一种做法是 epoch 前 10 个轮次只跑交叉熵后面的轮次再把 dice_loss 加进来。这样做的理由是让模型先用稳定梯度把整体结构学出来再加入区域重叠的精细约束。实操起来效果不错但没有严格的数学依据算是工程经验。3.3 数据增强对 loss 的连带影响用了随机翻转、随机裁剪这些增强后mask 和图片同步变换dice_loss 依然有效。但如果你用了 ColorJitter 调整亮度饱和度只有图片变换mask 不动——这在语义分割里是合理的因为颜色扰动不改变类别归属。需要注意的只有一个如果你做了随机裁剪裁剪后 mask 里可能完全没有前景像素这时候 dice_loss 的计算结果会偏大因为分母几乎只剩 smooth。应对方式是在 Dataset 里保证每个样本的 mask 前景面积不小于某个阈值比如裁剪后检查一下小于 5% 就重新裁剪或者直接跳过这个样本。4. 训练与 wandb 可视化监控 loss 曲线和分割结果4.1 基础训练流程的完整代码训练流程本身不复杂但有几个细节直接影响最终效果。学习率要分主干和分类头两部分设置——主干用预训练权重学习率乘 0.1分类头随机初始化用完整学习率。deeplabv3 内部的 BN 层在 batch size 比较小时最好冻结掉因为 BN 在小 batch 上统计的均值和方差漂移严重。实际项目中我见过很多次 batch size 2 训练出来的模型推理时边缘有一圈噪声就是 BN 没冻结导致的。from torch.optim import SGD from torch.optim.lr_scheduler import PolynomialLR model build_model(num_classes2).cuda() backbone_params [] classifier_params [] for name, param in model.named_parameters(): if classifier in name: classifier_params.append(param) else: backbone_params.append(param) optimizer SGD([ {params: backbone_params, lr: 0.001}, {params: classifier_params, lr: 0.01} ], momentum0.9, weight_decay1e-4) scheduler PolynomialLR(optimizer, total_iters50, power0.9) for epoch in range(50): model.train() running_loss 0.0 for images, masks in dataloader: images images.cuda() masks masks.cuda().long() optimizer.zero_grad() outputs model(images)[out] # outputs 形状是 (N, 2, H, W)需要对齐到 masks 的形状 loss_ce criterion(outputs, masks[:, 0]) pred_sigmoid torch.softmax(outputs, dim1)[:, 1, :, :] loss loss_ce 0.5 * dice_loss(pred_sigmoid, masks.float()) loss.backward() optimizer.step() running_loss loss.item() scheduler.step() print(fEpoch {epoch1}/50, Loss: {running_loss / len(dataloader):.4f})SGD 加 momentum 在这个任务上通常比 Adam 稳。Adam 收敛快但容易出现边界不干净的问题我见过不止一个项目在跑分割时用 Adam 训到最后每个物体边缘都像狗啃的换回 SGD 之后好了很多。如果你的训练数据本身就比较干净Adam 也可以接受但建议至少先用 SGD 跑一版当基线。4.2 wandb 怎么接进来wandbWeights Biases的作用不只是画 loss 曲线更能把每个 epoch 的分割结果图直接拉到网页上看。这比只看那串 loss 数字靠谱得多因为 loss 降了不代表分割边界对齐了——你需要在曲线之外看到预测 mask 长什么样。import wandb wandb.init(projectperson-seg-deeplabv3, config{ model: deeplabv3_resnet50, input_size: 320, batch_size: 8, epochs: 50, lr_backbone: 0.001, lr_classifier: 0.01, loss: cedice(0.5) }) # 训练循环内 for epoch in range(50): train_loss run_one_epoch() val_pred, val_mask predict_one_batch(val_loader, model) wandb.log({ train_loss: train_loss, epoch: epoch, val_prediction: wandb.Image(val_pred), val_mask: wandb.Image(val_mask) })wandb 的 log 不能放在每个 step 里要不网页会被刷爆。每轮 epoch 记录一次就够。第一次用的时候机器上需要先跑wandb login输入 API key之后训练脚本会同步到云端。如果你不想把数据传到外部也可以用wandb offline模式日志先存在本地后面一次性同步。这套组合的效果就是模型训练 50 轮之后你打开 wandb 的页面能看到每条曲线的走向、每个 epoch 的预测 mask 和真实 mask 并排对比哪一轮开始过拟合一目了然。4.3 训练轮次与收敛判断50 轮是一个比较常见的配置但实际收敛点取决于数据量和任务难度。数据量几百张图的话30 轮左右 loss 就基本平稳了。判断标准不是 loss 绝对值而是 val 部分的 mIoU 是否还在上升——如果连续 10 轮 mIoU 都没超过历史最好值就可以早停。这个早停逻辑我强烈建议写到脚本里因为训练中人工盯着到底什么时候停很容易判断失误。5. 常见问题排查显存溢出、全黑预测与 dice 为 0 的四个坑5.1 显存溢出ResNet50 backbone 比想象中吃显存现象训练脚本刚跑到第一个 step直接报 CUDA out of memory或者跑了几个 step 之后在 random 位置崩掉。原因deeplabv3_resnet50 在 320x320 输入下batch size 8 的显存需求接近 7G。如果你用的是 6G 显存的卡batch size 6 都可能悬。解决先降 batch size 到 2 或 4同时把输入尺寸从 320 降到 256。还可以torch.cuda.empty_cache()在每轮 epoch 结束后清一下碎显存——PyTorch 的显存分配器有时候不会自动释放所有缓存块。torch.cuda.empty_cache() # 在每个 epoch 结束时调用5.2 预测结果全黑softmax 后取错了类别通道现象模型训练完了预测图输出全黑或者全是白色。原因deeplabv3 输出的 logits 形状是 (N, 2, H, W)channel 0 是背景、channel 1 是前景。预测时如果直接取了 channel 0得到的就是背景图。这类问题在写预测代码的当天很容易犯。解决明确取索引 1 作为前景并 squeeze 掉多余的通道。output model(img_tensor)[out] pred torch.argmax(output, dim1).squeeze(0) # (H, W)像素值为 0 或 1 # 如果想保存成黑白图需要乘 255 再转 uint85.3 dice_loss 一直是 0 或恒定值mask 没归一化现象训练日志里 dice_loss 从第一个 epoch 开始就是 0或者一直恒定在一个固定值。原因如果 mask 图像是用 PIL 读入后直接转 tensor像素值是 0 到 255 而不是 0 到 1。mask 全黑时 target 为 0 的像素占了 99%dice 系数天然接近 0。解决加载 mask 后必须除以 255 归一化同时转成 float。mask_tensor torch.from_numpy(np.array(mask)).float().unsqueeze(0) / 255.05.4 wandb 网页上看不到图片log 里传的格式不对现象wandb 的 run 能建立loss 曲线有数据但 Val Prediction 区域空白或报错。原因wandb.Image 接受 numpy 数组或者 PIL Image不接受 torch tensor。很多人直接把 tensor 传进去就完事了结果 wandb 打不开。解决先用torchvision.utils.make_grid把 batch 拼成一张大图再.cpu().numpy()转成 numpy最后传给 wandb.Image。from torchvision.utils import make_grid grid make_grid(pred_img_tensor, normalizeTrue, value_range(0, 1)) wandb.log({val_prediction: wandb.Image(grid.permute(1, 2, 0).cpu().numpy())})5.5 同一份数据训练精度波动大没有固定随机种子现象同一份数据同一个脚本两次训练出来的 mIoU 差别超过 5 个点。原因PyTorch 的 DataLoader 有随机 shuffle模型初始化有随机性。如果你在对比不同模型的优劣这个波动会直接影响结论。解决训练脚本开头固定三个随机种子。import random import numpy as np torch.manual_seed(42) random.seed(42) np.random.seed(42)6. 预测推理与落地技巧从单张图片到批量处理训练完了最后落到使用场景上。这个 zip 里已经有一张 output.jpg说明作者给你示范了单张预测的效果。实际工程里更常见的是对一个目录下几百张图片批量抠图保存成透明 PNG 或者贴到新背景上。批量预测要注意两点一是模型要切到 eval 模式并关闭梯度计算二是输入图片做和训练时一致的预处理包括尺寸缩放和归一化。from torchvision import transforms from PIL import Image import numpy as np model.eval() transform transforms.Compose([ transforms.Resize((320, 320)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) def predict_single(model, img_path, save_path): img Image.open(img_path).convert(RGB) w, h img.size tensor transform(img).unsqueeze(0).cuda() with torch.no_grad(): output model(tensor)[out] pred torch.softmax(output, dim1)[0, 1].cpu().numpy() # 二值化0.5 作为前景背景阈值 mask (pred 0.5).astype(np.uint8) * 255 mask_img Image.fromarray(mask).resize((w, h), Image.NEAREST) # 拼接原图和 mask 并排保存方便人工检查 bg np.zeros((h, w, 4), dtypenp.uint8) bg[:, :, :3] np.array(img) bg[:, :, 3] mask Image.fromarray(bg).save(save_path)这段代码最后输出的是带透明通道的 RGBA 图背景是黑色人物位置保留原色。如果只想拿 mask直接保存 mask_img。看到这里可能有人注意到 tensor 归一化的 mean/std 用的是 ImageNet 的标准因为 deeplabv3 预训练就是在这个分布上做的不要改成你自己的均值。训练时如果也用了同样的预处理那预测既然要和训练对齐就继续用它。抠图质量卡的其实是边缘而边缘质量又由模型的输出分辨率决定。deeplabv3 输出的 mask 是 320x320直接 resize 回原图尺寸时锯齿会很明显。我一般会在二值化前把 pred 概率图先做一个 2x 的上采样比如用双线性插值把 (320,320) 放大到 (640,640)这时候原来一个像素的判断变成了 2x2 区域再二值化边缘会柔和一些。更讲究的做法是用引导滤波或 CRF 做后处理但那个就属于加分项了。有几条实操习惯是做完这个项目后养成的。每次训练前固定随机种子调完 batch size 先跑两步验证 loss 能正常回传再放全量训练每个模型训练完顺手导出几张验证集预测图放到 wandb 上避免被单次 loss 数值欺骗。从那以后每一次在这个资源上改动模型结构或者换主干我都会强制走一遍这四步——固定种子、跑两个 mini-batch 验证、训练中盯 wandb 图、预测前确认类别通道。这套流程看着笨但能防住大多数和数据集无关的愚蠢错误。希望帮到你。本文还有配套的精品资源点击获取