简介本资源面向具备PyTorch基础的深度学习学习者与算法工程师聚焦二分类语义分割的人物抠图任务提供一套可复现的完整工程代码与数据集。包内共约2000个文件以png图像数据为主辅以py训练与预测脚本、pyc缓存、Dockerfile、sh启动脚本及README说明文档压缩包约990.97MB目录结构清晰便于按模块查阅与调试。内容围绕segmentation_models.pytorch框架展开涵盖语义分割数据集的加载与使用、交叉熵与DiceLoss组合损失、wandb训练可视化以及二分类分割的训练与预测全流程帮助读者理解常用做法并快速复现算法。已有1764人学习下载适合希望掌握人物抠图实战、对照代码查漏补缺的中高级开发者参考。1. 用 segmentation_models.pytorch 做人物抠图为什么它比直接跑 UNet 更省事人物抠图这件事很多人第一反应是找一套 UNet 源码改改输入输出通道就开始训。但真跑起来会发现编码器权重从零训、数据增强要自己搭、学习率调度要自己调光是把 loss 降到能看的水平就得折腾好几天。segmentation_models.pytorch后面简称 smp解决的正是这个前置成本问题它把主流编码器ResNet、EfficientNet、MobileNet 等和分割头做了标准化封装换骨干网络只需要改一个字符串参数预训练权重直接加载省掉了大量重复劳动。这篇文章面向的是想用 smp 做人物抠图但还没跑通完整流程的人。不管你是要做证件照换底、直播虚拟背景还是电商商品图的人物分离核心链路是一样的数据准备、模型构建、训练调参、推理后处理。我会按这条链路把每一步的参数含义和踩坑点讲清楚代码可以直接抄。地毯图像分割系统、广告牌图像分割系统这类项目里的人物区域提取用的也是同一套逻辑只是数据集和后处理策略不同。2. 人物抠图的数据准备从标注格式到 Dataset 怎么写2.1 抠图任务的数据长什么样人物抠图本质是二分类语义分割前景人为 1背景为 0。输入是 RGB 三通道图像标签是单通道掩码图mask像素值只有 0 和 255或 0 和 1。常见的数据来源有三种一是公开数据集如 Supervisely Person Dataset、AISegment二是用标注工具LabelMe、CVAT自己标三是用 matting 算法生成粗掩码再人工修正。标注格式方面LabelMe 导出的是 JSON需要转成 PNG 掩码CVAT 可以直接导出 mask。不管哪种来源最终统一成这样的目录结构dataset/ images/ 001.jpg 002.jpg masks/ 001.png 002.png图像和掩码文件名必须一一对应尺寸也要一致。如果原始掩码是灰度图读取后要二值化处理否则训练时 loss 会出问题。2.2 Dataset 和 DataLoader 的完整写法smp 本身不提供 Dataset 类需要自己写。下面是一个可直接用的实现import os import cv2 import numpy as np import torch from torch.utils.data import Dataset, DataLoader import albumentations as A from albumentations.pytorch import ToTensorV2 class PersonSegDataset(Dataset): def __init__(self, img_dir, mask_dir, transformNone): self.img_dir img_dir self.mask_dir mask_dir self.transform transform # 只保留 images 目录下真实存在的文件避免残留文件导致报错 self.names [f for f in os.listdir(img_dir) if f.lower().endswith((.jpg, .jpeg, .png))] def __len__(self): return len(self.names) def __getitem__(self, idx): name self.names[idx] img_path os.path.join(self.img_dir, name) # 掩码统一用 png文件名 stem 与图像一致 mask_name os.path.splitext(name)[0] .png mask_path os.path.join(self.mask_dir, mask_name) image cv2.imread(img_path) image cv2.cvtColor(image, cv2.COLOR_BGR2RGB) mask cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) # 二值化大于 127 视为前景统一成 0/1 的 float mask (mask 127).astype(np.float32) if self.transform: augmented self.transform(imageimage, maskmask) image augmented[image] mask augmented[mask] # 掩码增加通道维度 - (1, H, W) mask mask.unsqueeze(0) return image, mask # 训练集增强几何变换 颜色抖动注意 mask 只做几何变换 train_transform A.Compose([ A.Resize(320, 320), A.HorizontalFlip(p0.5), A.ShiftScaleRotate(shift_limit0.1, scale_limit0.15, rotate_limit15, p0.5), A.RandomBrightnessContrast(p0.3), A.Normalize(mean(0.485, 0.456, 0.406), std(0.229, 0.224, 0.225)), ToTensorV2(), ]) # 验证集只做 resize 和归一化不做随机增强 val_transform A.Compose([ A.Resize(320, 320), A.Normalize(mean(0.485, 0.456, 0.406), std(0.229, 0.224, 0.225)), ToTensorV2(), ])这段代码有几个关键点需要说明。第一掩码二值化阈值选 127 是因为标注工具导出的掩码边缘可能有抗锯齿产生的灰度值用 127 做分界能保留主要前景区域。第二albumentations 的 Compose 对 image 和 mask 同时做几何变换保证空间对齐但颜色变换只作用于 image不会影响 mask。第三Normalize 的均值和标准差用的是 ImageNet 统计值因为 smp 的编码器预训练权重来自 ImageNet输入分布要匹配。DataLoader 的配置train_dataset PersonSegDataset(dataset/images, dataset/masks, train_transform) train_loader DataLoader( train_dataset, batch_size8, # 显存 8G 以下建议降到 4 shuffleTrue, num_workers4, # Windows 下设为 0否则可能卡死 pin_memoryTrue, # GPU 训练时开启加速数据传输 drop_lastTrue # 丢弃最后一个不完整 batch避免 BN 层报错 )num_workers在 Windows 上是个经典坑设大于 0 时如果 Dataset 里有 cv2 读取操作可能直接卡住不动。drop_lastTrue是为了避免最后一个 batch 只有 1 张图时 BatchNorm 层计算出问题。3. 用 smp 构建分割模型编码器选型与解码器参数怎么定3.1 编码器选型的三个实际考量smp 支持的编码器有几十种人物抠图场景下我一般从这三个维度选编码器参数量推理速度适合场景resnet34中中精度和速度平衡首选基线efficientnet-b3中高中慢精度要求高显存充足mobilenet_v2低快移动端部署或实时推理mit_b0低中Transformer 结构边缘细节好新手建议从resnet34起步跑通全流程后再换。如果目标是移动端部署直接上mobilenet_v2但要注意它的精度上限比 ResNet 低一截后处理需要多花功夫。3.2 模型构建代码与参数说明import segmentation_models_pytorch as smp model smp.Unet( encoder_nameresnet34, # 编码器名称换成其他只需改这里 encoder_weightsimagenet, # 预训练权重首次运行会自动下载 in_channels3, # 输入 RGB 三通道 classes1, # 二分类输出 1 个通道 activationNone, # 训练时输出 logits推理时再 sigmoid decoder_attention_typescse, # 通道注意力对小目标边缘有提升 )activationNone是训练阶段的关键设置。smp 的 loss 函数如 BCEWithLogitsLoss、DiceLoss内部会做 sigmoid如果模型输出已经过 sigmoid会导致梯度消失。推理时再手动加 sigmoid 或 softmax。decoder_attention_typescse是空间和通道注意力模块对人物边缘头发丝、手指缝隙的分割精度有可见提升代价是推理速度略降。如果追求极致速度可以去掉。3.3 损失函数与优化器的搭配人物抠图常见的问题是前景背景像素不平衡——人物通常只占图像的 20%40%。纯 BCE 会让模型倾向于预测全背景。我一般用 Dice BCE 组合import torch import torch.nn as nn dice_loss smp.losses.DiceLoss(modebinary) bce_loss nn.BCEWithLogitsLoss() def criterion(pred, target): # Dice 关注区域重叠度BCE 关注逐像素分类 return dice_loss(pred, target) bce_loss(pred, target) optimizer torch.optim.AdamW( model.parameters(), lr1e-3, # 初始学习率smp 官方示例常用 1e-3 weight_decay1e-4 # 权重衰减防止过拟合 ) scheduler torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max50, eta_min1e-6 )Dice Loss 对类别不平衡天然鲁棒BCE 提供稳定的逐像素梯度。两者相加是分割任务里经过验证的组合。学习率用余弦退火从 1e-3 降到 1e-6比固定学习率收敛更稳。4. 训练循环与验证loss 不降时该看哪些指标4.1 训练循环的完整实现import torch from tqdm import tqdm device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) def train_one_epoch(model, loader, optimizer, criterion, device): model.train() total_loss 0 for images, masks in tqdm(loader, desctrain): images images.to(device) masks masks.to(device) optimizer.zero_grad() preds model(images) # 输出 logits形状 (B, 1, H, W) loss criterion(preds, masks) loss.backward() optimizer.step() total_loss loss.item() return total_loss / len(loader) torch.no_grad() def validate(model, loader, criterion, device): model.eval() total_loss 0 total_iou 0 for images, masks in tqdm(loader, descval): images images.to(device) masks masks.to(device) preds model(images) loss criterion(preds, masks) total_loss loss.item() # 计算 IoU先 sigmoid 再阈值化 probs torch.sigmoid(preds) pred_bin (probs 0.5).float() intersection (pred_bin * masks).sum() union pred_bin.sum() masks.sum() - intersection iou (intersection 1e-6) / (union 1e-6) total_iou iou.item() return total_loss / len(loader), total_iou / len(loader)验证阶段计算 IoU 比只看 loss 更直观。IoU 低于 0.7 说明模型还没学到有效特征高于 0.9 基本可用。注意torch.no_grad()一定要加否则验证会占显存且拖慢速度。4.2 loss 不降时的排查顺序训练中 loss 卡住不动按这个顺序查第一检查掩码是否正确加载。把 Dataset 里读出的 mask 可视化一下确认前景是白色、背景是黑色而不是全黑或全白。标注格式错误是最常见的原因。第二检查学习率是否过大。loss 震荡不降甚至上升把 lr 降到 1e-4 再试。smp 的预训练编码器对学习率比较敏感1e-3 有时偏大。第三检查 loss 函数和模型输出是否匹配。模型activationNone输出 logitsloss 必须用带 logits 的版本BCEWithLogitsLoss不能先用 sigmoid 再算 BCE。第四检查数据增强是否过猛。ShiftScaleRotate 的 scale_limit 设到 0.3 以上时人物可能被裁出画面导致掩码全黑。先把增强关掉跑一轮确认模型能过拟合小数据集再逐步加增强。4.3 训练轮次与保存策略best_iou 0 for epoch in range(50): train_loss train_one_epoch(model, train_loader, optimizer, criterion, device) val_loss, val_iou validate(model, val_loader, criterion, device) scheduler.step() print(fEpoch {epoch}: train_loss{train_loss:.4f}, fval_loss{val_loss:.4f}, val_iou{val_iou:.4f}) # 只保存验证 IoU 最高的模型避免保存过拟合的后期权重 if val_iou best_iou: best_iou val_iou torch.save(model.state_dict(), best_model.pth)保存策略用「验证集最优」而不是「最后一轮」是因为分割任务通常在 2030 轮后开始过拟合最后一轮的权重反而不是最好的。5. 推理与后处理从 logits 到可用掩码的完整链路5.1 单张图像推理代码import cv2 import numpy as np import torch def predict_mask(model, image_path, device, input_size320): model.eval() # 读取并预处理 image cv2.imread(image_path) image cv2.cvtColor(image, cv2.COLOR_BGR2RGB) h, w image.shape[:2] # 保存原始尺寸用于还原 resized cv2.resize(image, (input_size, input_size)) tensor resized.astype(np.float32) / 255.0 tensor (tensor - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225] tensor torch.from_numpy(tensor).permute(2, 0, 1).unsqueeze(0).float() tensor tensor.to(device) with torch.no_grad(): logits model(tensor) prob torch.sigmoid(logits) # 转成 0~1 概率 prob prob.squeeze().cpu().numpy() # (H, W) # 还原到原始尺寸 prob cv2.resize(prob, (w, h)) return prob def refine_mask(prob, threshold0.5, kernel_size5): # 阈值化 mask (prob threshold).astype(np.uint8) * 255 # 形态学开运算去噪点闭运算填小洞 kernel cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)) mask cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel) mask cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel) # 保留最大连通域去掉误检的小区域 contours, _ cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) if contours: largest max(contours, keycv2.contourArea) clean np.zeros_like(mask) cv2.drawContours(clean, [largest], -1, 255, -1) mask clean return mask推理链路的关键在于模型输出的是 320×320 的 logits要 resize 回原图尺寸再做后处理。直接在小尺寸上做形态学操作再放大边缘会糊。5.2 边缘优化的两个实用技巧第一个是阈值调整。默认 0.5 对大多数场景够用但如果发现人物边缘缺了一块比如深色衣服和深色背景交界处把阈值降到 0.30.4 能找回更多前景像素。反之如果背景有误检提高到 0.6。第二个是引导滤波guided filter。用原图作为引导图对概率图做滤波能让边缘贴合原始图像的纹理边界def guided_filter_mask(prob, image, radius8, eps1e-4): # prob: 0~1 浮点概率图, image: 原始 RGB 图 guide image.astype(np.float32) / 255.0 prob_3c np.stack([prob] * 3, axis-1) # 简化版引导滤波用均值滤波近似 mean_guide cv2.blur(guide, (radius, radius)) mean_prob cv2.blur(prob_3c, (radius, radius)) corr_guide cv2.blur(guide * prob_3c, (radius, radius)) cov corr_guide - mean_guide * mean_prob var cv2.blur(guide * guide, (radius, radius)) - mean_guide ** 2 a cov / (var eps) b mean_prob - a * mean_guide mean_a cv2.blur(a, (radius, radius)) mean_b cv2.blur(b, (radius, radius)) refined mean_a * guide mean_b return refined.mean(axis-1)引导滤波的radius控制边缘平滑范围eps控制边缘保留强度。人物抠图场景下 radius8、eps1e-4 是比较稳的起点。5.3 抠图输出与背景合成拿到最终掩码后合成新背景def composite(image_path, mask, bg_color(255, 255, 255)): image cv2.imread(image_path) mask_3c cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR) / 255.0 bg np.full_like(image, bg_color, dtypenp.uint8) result (image * mask_3c bg * (1 - mask_3c)).astype(np.uint8) return result如果要换的是图片背景而非纯色把bg换成 resize 到同尺寸的背景图即可。注意掩码边缘可以做 12 像素的高斯模糊合成后过渡更自然不会有锯齿感。6. 避坑与排查人物抠图训练中最容易翻车的 5 个地方现象一训练 loss 正常下降但推理时掩码全白或全黑。原因训练时用了activationNone推理时忘了加 sigmoid直接把 logits 当概率用。logits 范围是负无穷到正无穷阈值化后自然全白或全黑。 解决推理代码里确认有torch.sigmoid(logits)这一步或者加载模型时设activationsigmoid但训练时必须为 None。现象二验证集 IoU 很高但实际图片抠出来边缘毛糙、有大量小碎块。原因验证集和训练集来自同一分布且验证时只做了 resize没有覆盖真实场景的多样性。另外模型输出的概率图本身就有噪声。 解决后处理加形态学开闭运算和最大连通域保留。如果碎块仍然多检查训练数据里是否有标注质量差的样本把 IoU 低于 0.5 的标注挑出来修正或丢弃。现象三训练到 10 轮左右 loss 突然变成 NaN。原因学习率过大导致梯度爆炸或者 Dice Loss 在预测全零时出现除零。另外混合精度训练AMP下 BCE 有时会溢出。 解决先把学习率降到 1e-4加梯度裁剪torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)。Dice Loss 里加smooth1.0参数避免除零。如果用了 AMP把 loss 计算部分强制转成 float32。现象四换编码器后显存不够batch_size 降到 2 还是 OOM。原因不同编码器的参数量和中间特征图大小差异很大。efficientnet-b5 比 resnet34 的显存占用高好几倍。 解决先降输入分辨率320→256再考虑换更小的编码器。另外检查是否在验证时忘了加torch.no_grad()这会导致验证阶段也构建计算图显存翻倍。现象五人物边缘的头发丝区域总是被切成背景。原因头发丝区域像素占比极小Dice Loss 对这种细结构不敏感。另外下采样过程中细节信息丢失严重。 解决换用decoder_attention_typescse增强细节关注损失函数里加边界加权对掩码边缘像素给更高权重或者换mit_b0这类 Transformer 编码器它对长距离依赖和细结构的分割效果更好。如果还不行考虑用 matting 方法做后处理精修。7. 把抠图精度再推一步从分割到 matting 的衔接思路分割模型输出的掩码是硬边界的阈值化之后只有 0 和 255 两个值。但真实的人物边缘——尤其是头发、半透明衣物、运动模糊区域——需要 alpha 渐变才能自然合成。这就是 matting 要解决的问题。一个实用的衔接方案是用 smp 分割模型生成粗掩码然后在掩码边缘的过渡带膨胀掩码减去腐蚀掩码得到的区域内用 KNN matting 或 Closed-form matting 做精修。这样既利用了分割模型的语义理解能力又用 matting 补上了边缘细节。def trimap_from_mask(mask, erode_size5, dilate_size10): # 生成三分图确定前景255确定背景0未知区域128 kernel_e cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (erode_size, erode_size)) kernel_d cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (dilate_size, dilate_size)) fg cv2.erode(mask, kernel_e, iterations1) # 腐蚀得到确定前景 bg cv2.dilate(mask, kernel_d, iterations1) # 膨胀得到确定背景 trimap np.full_like(mask, 128) # 默认未知 trimap[bg 0] 0 # 确定背景 trimap[fg 255] 255 # 确定前景 return trimap拿到 trimap 后可以用pymatting库做 alpha 估计from pymatting import estimate_alpha_cf def matting_refine(image, trimap): # image 归一化到 0~1 img_float image.astype(np.float64) / 255.0 alpha estimate_alpha_cf(img_float, trimap) return (alpha * 255).astype(np.uint8)estimate_alpha_cf是 Closed-form matting 的实现对头发丝区域效果明显。代价是速度慢一张 1080P 图大概要几秒。如果做批量处理可以只对 trimap 的未知区域做计算确定区域直接赋值能省不少时间。验证 matting 效果的方法把 alpha 掩码和原图做合成换一个纯色背景放大看边缘过渡是否自然。如果边缘有白边或黑边说明 trimap 的腐蚀/膨胀参数需要调——白边通常是确定前景范围太小黑边是确定背景范围太小。我自己做人物抠图项目时习惯在训练完分割模型后先跑一批测试图把 IoU 低于 0.85 的样本单独拎出来看。十有八九是标注问题或者极端光照条件而不是模型本身不行。把这几张图的标注修一修重新训一轮比调参折腾半天管用得多。希望帮到你。本文还有配套的精品资源点击获取
简介本资源面向具备PyTorch基础的深度学习学习者与算法工程师聚焦二分类语义分割的人物抠图任务提供一套可复现的完整工程代码与数据集。包内共约2000个文件以png图像数据为主辅以py训练与预测脚本、pyc缓存、Dockerfile、sh启动脚本及README说明文档压缩包约990.97MB目录结构清晰便于按模块查阅与调试。内容围绕segmentation_models.pytorch框架展开涵盖语义分割数据集的加载与使用、交叉熵与DiceLoss组合损失、wandb训练可视化以及二分类分割的训练与预测全流程帮助读者理解常用做法并快速复现算法。已有1764人学习下载适合希望掌握人物抠图实战、对照代码查漏补缺的中高级开发者参考。1. 用 segmentation_models.pytorch 做人物抠图为什么它比直接跑 UNet 更省事人物抠图这件事很多人第一反应是找一套 UNet 源码改改输入输出通道就开始训。但真跑起来会发现编码器权重从零训、数据增强要自己搭、学习率调度要自己调光是把 loss 降到能看的水平就得折腾好几天。segmentation_models.pytorch后面简称 smp解决的正是这个前置成本问题它把主流编码器ResNet、EfficientNet、MobileNet 等和分割头做了标准化封装换骨干网络只需要改一个字符串参数预训练权重直接加载省掉了大量重复劳动。这篇文章面向的是想用 smp 做人物抠图但还没跑通完整流程的人。不管你是要做证件照换底、直播虚拟背景还是电商商品图的人物分离核心链路是一样的数据准备、模型构建、训练调参、推理后处理。我会按这条链路把每一步的参数含义和踩坑点讲清楚代码可以直接抄。地毯图像分割系统、广告牌图像分割系统这类项目里的人物区域提取用的也是同一套逻辑只是数据集和后处理策略不同。2. 人物抠图的数据准备从标注格式到 Dataset 怎么写2.1 抠图任务的数据长什么样人物抠图本质是二分类语义分割前景人为 1背景为 0。输入是 RGB 三通道图像标签是单通道掩码图mask像素值只有 0 和 255或 0 和 1。常见的数据来源有三种一是公开数据集如 Supervisely Person Dataset、AISegment二是用标注工具LabelMe、CVAT自己标三是用 matting 算法生成粗掩码再人工修正。标注格式方面LabelMe 导出的是 JSON需要转成 PNG 掩码CVAT 可以直接导出 mask。不管哪种来源最终统一成这样的目录结构dataset/ images/ 001.jpg 002.jpg masks/ 001.png 002.png图像和掩码文件名必须一一对应尺寸也要一致。如果原始掩码是灰度图读取后要二值化处理否则训练时 loss 会出问题。2.2 Dataset 和 DataLoader 的完整写法smp 本身不提供 Dataset 类需要自己写。下面是一个可直接用的实现import os import cv2 import numpy as np import torch from torch.utils.data import Dataset, DataLoader import albumentations as A from albumentations.pytorch import ToTensorV2 class PersonSegDataset(Dataset): def __init__(self, img_dir, mask_dir, transformNone): self.img_dir img_dir self.mask_dir mask_dir self.transform transform # 只保留 images 目录下真实存在的文件避免残留文件导致报错 self.names [f for f in os.listdir(img_dir) if f.lower().endswith((.jpg, .jpeg, .png))] def __len__(self): return len(self.names) def __getitem__(self, idx): name self.names[idx] img_path os.path.join(self.img_dir, name) # 掩码统一用 png文件名 stem 与图像一致 mask_name os.path.splitext(name)[0] .png mask_path os.path.join(self.mask_dir, mask_name) image cv2.imread(img_path) image cv2.cvtColor(image, cv2.COLOR_BGR2RGB) mask cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) # 二值化大于 127 视为前景统一成 0/1 的 float mask (mask 127).astype(np.float32) if self.transform: augmented self.transform(imageimage, maskmask) image augmented[image] mask augmented[mask] # 掩码增加通道维度 - (1, H, W) mask mask.unsqueeze(0) return image, mask # 训练集增强几何变换 颜色抖动注意 mask 只做几何变换 train_transform A.Compose([ A.Resize(320, 320), A.HorizontalFlip(p0.5), A.ShiftScaleRotate(shift_limit0.1, scale_limit0.15, rotate_limit15, p0.5), A.RandomBrightnessContrast(p0.3), A.Normalize(mean(0.485, 0.456, 0.406), std(0.229, 0.224, 0.225)), ToTensorV2(), ]) # 验证集只做 resize 和归一化不做随机增强 val_transform A.Compose([ A.Resize(320, 320), A.Normalize(mean(0.485, 0.456, 0.406), std(0.229, 0.224, 0.225)), ToTensorV2(), ])这段代码有几个关键点需要说明。第一掩码二值化阈值选 127 是因为标注工具导出的掩码边缘可能有抗锯齿产生的灰度值用 127 做分界能保留主要前景区域。第二albumentations 的 Compose 对 image 和 mask 同时做几何变换保证空间对齐但颜色变换只作用于 image不会影响 mask。第三Normalize 的均值和标准差用的是 ImageNet 统计值因为 smp 的编码器预训练权重来自 ImageNet输入分布要匹配。DataLoader 的配置train_dataset PersonSegDataset(dataset/images, dataset/masks, train_transform) train_loader DataLoader( train_dataset, batch_size8, # 显存 8G 以下建议降到 4 shuffleTrue, num_workers4, # Windows 下设为 0否则可能卡死 pin_memoryTrue, # GPU 训练时开启加速数据传输 drop_lastTrue # 丢弃最后一个不完整 batch避免 BN 层报错 )num_workers在 Windows 上是个经典坑设大于 0 时如果 Dataset 里有 cv2 读取操作可能直接卡住不动。drop_lastTrue是为了避免最后一个 batch 只有 1 张图时 BatchNorm 层计算出问题。3. 用 smp 构建分割模型编码器选型与解码器参数怎么定3.1 编码器选型的三个实际考量smp 支持的编码器有几十种人物抠图场景下我一般从这三个维度选编码器参数量推理速度适合场景resnet34中中精度和速度平衡首选基线efficientnet-b3中高中慢精度要求高显存充足mobilenet_v2低快移动端部署或实时推理mit_b0低中Transformer 结构边缘细节好新手建议从resnet34起步跑通全流程后再换。如果目标是移动端部署直接上mobilenet_v2但要注意它的精度上限比 ResNet 低一截后处理需要多花功夫。3.2 模型构建代码与参数说明import segmentation_models_pytorch as smp model smp.Unet( encoder_nameresnet34, # 编码器名称换成其他只需改这里 encoder_weightsimagenet, # 预训练权重首次运行会自动下载 in_channels3, # 输入 RGB 三通道 classes1, # 二分类输出 1 个通道 activationNone, # 训练时输出 logits推理时再 sigmoid decoder_attention_typescse, # 通道注意力对小目标边缘有提升 )activationNone是训练阶段的关键设置。smp 的 loss 函数如 BCEWithLogitsLoss、DiceLoss内部会做 sigmoid如果模型输出已经过 sigmoid会导致梯度消失。推理时再手动加 sigmoid 或 softmax。decoder_attention_typescse是空间和通道注意力模块对人物边缘头发丝、手指缝隙的分割精度有可见提升代价是推理速度略降。如果追求极致速度可以去掉。3.3 损失函数与优化器的搭配人物抠图常见的问题是前景背景像素不平衡——人物通常只占图像的 20%40%。纯 BCE 会让模型倾向于预测全背景。我一般用 Dice BCE 组合import torch import torch.nn as nn dice_loss smp.losses.DiceLoss(modebinary) bce_loss nn.BCEWithLogitsLoss() def criterion(pred, target): # Dice 关注区域重叠度BCE 关注逐像素分类 return dice_loss(pred, target) bce_loss(pred, target) optimizer torch.optim.AdamW( model.parameters(), lr1e-3, # 初始学习率smp 官方示例常用 1e-3 weight_decay1e-4 # 权重衰减防止过拟合 ) scheduler torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max50, eta_min1e-6 )Dice Loss 对类别不平衡天然鲁棒BCE 提供稳定的逐像素梯度。两者相加是分割任务里经过验证的组合。学习率用余弦退火从 1e-3 降到 1e-6比固定学习率收敛更稳。4. 训练循环与验证loss 不降时该看哪些指标4.1 训练循环的完整实现import torch from tqdm import tqdm device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) def train_one_epoch(model, loader, optimizer, criterion, device): model.train() total_loss 0 for images, masks in tqdm(loader, desctrain): images images.to(device) masks masks.to(device) optimizer.zero_grad() preds model(images) # 输出 logits形状 (B, 1, H, W) loss criterion(preds, masks) loss.backward() optimizer.step() total_loss loss.item() return total_loss / len(loader) torch.no_grad() def validate(model, loader, criterion, device): model.eval() total_loss 0 total_iou 0 for images, masks in tqdm(loader, descval): images images.to(device) masks masks.to(device) preds model(images) loss criterion(preds, masks) total_loss loss.item() # 计算 IoU先 sigmoid 再阈值化 probs torch.sigmoid(preds) pred_bin (probs 0.5).float() intersection (pred_bin * masks).sum() union pred_bin.sum() masks.sum() - intersection iou (intersection 1e-6) / (union 1e-6) total_iou iou.item() return total_loss / len(loader), total_iou / len(loader)验证阶段计算 IoU 比只看 loss 更直观。IoU 低于 0.7 说明模型还没学到有效特征高于 0.9 基本可用。注意torch.no_grad()一定要加否则验证会占显存且拖慢速度。4.2 loss 不降时的排查顺序训练中 loss 卡住不动按这个顺序查第一检查掩码是否正确加载。把 Dataset 里读出的 mask 可视化一下确认前景是白色、背景是黑色而不是全黑或全白。标注格式错误是最常见的原因。第二检查学习率是否过大。loss 震荡不降甚至上升把 lr 降到 1e-4 再试。smp 的预训练编码器对学习率比较敏感1e-3 有时偏大。第三检查 loss 函数和模型输出是否匹配。模型activationNone输出 logitsloss 必须用带 logits 的版本BCEWithLogitsLoss不能先用 sigmoid 再算 BCE。第四检查数据增强是否过猛。ShiftScaleRotate 的 scale_limit 设到 0.3 以上时人物可能被裁出画面导致掩码全黑。先把增强关掉跑一轮确认模型能过拟合小数据集再逐步加增强。4.3 训练轮次与保存策略best_iou 0 for epoch in range(50): train_loss train_one_epoch(model, train_loader, optimizer, criterion, device) val_loss, val_iou validate(model, val_loader, criterion, device) scheduler.step() print(fEpoch {epoch}: train_loss{train_loss:.4f}, fval_loss{val_loss:.4f}, val_iou{val_iou:.4f}) # 只保存验证 IoU 最高的模型避免保存过拟合的后期权重 if val_iou best_iou: best_iou val_iou torch.save(model.state_dict(), best_model.pth)保存策略用「验证集最优」而不是「最后一轮」是因为分割任务通常在 2030 轮后开始过拟合最后一轮的权重反而不是最好的。5. 推理与后处理从 logits 到可用掩码的完整链路5.1 单张图像推理代码import cv2 import numpy as np import torch def predict_mask(model, image_path, device, input_size320): model.eval() # 读取并预处理 image cv2.imread(image_path) image cv2.cvtColor(image, cv2.COLOR_BGR2RGB) h, w image.shape[:2] # 保存原始尺寸用于还原 resized cv2.resize(image, (input_size, input_size)) tensor resized.astype(np.float32) / 255.0 tensor (tensor - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225] tensor torch.from_numpy(tensor).permute(2, 0, 1).unsqueeze(0).float() tensor tensor.to(device) with torch.no_grad(): logits model(tensor) prob torch.sigmoid(logits) # 转成 0~1 概率 prob prob.squeeze().cpu().numpy() # (H, W) # 还原到原始尺寸 prob cv2.resize(prob, (w, h)) return prob def refine_mask(prob, threshold0.5, kernel_size5): # 阈值化 mask (prob threshold).astype(np.uint8) * 255 # 形态学开运算去噪点闭运算填小洞 kernel cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)) mask cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel) mask cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel) # 保留最大连通域去掉误检的小区域 contours, _ cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) if contours: largest max(contours, keycv2.contourArea) clean np.zeros_like(mask) cv2.drawContours(clean, [largest], -1, 255, -1) mask clean return mask推理链路的关键在于模型输出的是 320×320 的 logits要 resize 回原图尺寸再做后处理。直接在小尺寸上做形态学操作再放大边缘会糊。5.2 边缘优化的两个实用技巧第一个是阈值调整。默认 0.5 对大多数场景够用但如果发现人物边缘缺了一块比如深色衣服和深色背景交界处把阈值降到 0.30.4 能找回更多前景像素。反之如果背景有误检提高到 0.6。第二个是引导滤波guided filter。用原图作为引导图对概率图做滤波能让边缘贴合原始图像的纹理边界def guided_filter_mask(prob, image, radius8, eps1e-4): # prob: 0~1 浮点概率图, image: 原始 RGB 图 guide image.astype(np.float32) / 255.0 prob_3c np.stack([prob] * 3, axis-1) # 简化版引导滤波用均值滤波近似 mean_guide cv2.blur(guide, (radius, radius)) mean_prob cv2.blur(prob_3c, (radius, radius)) corr_guide cv2.blur(guide * prob_3c, (radius, radius)) cov corr_guide - mean_guide * mean_prob var cv2.blur(guide * guide, (radius, radius)) - mean_guide ** 2 a cov / (var eps) b mean_prob - a * mean_guide mean_a cv2.blur(a, (radius, radius)) mean_b cv2.blur(b, (radius, radius)) refined mean_a * guide mean_b return refined.mean(axis-1)引导滤波的radius控制边缘平滑范围eps控制边缘保留强度。人物抠图场景下 radius8、eps1e-4 是比较稳的起点。5.3 抠图输出与背景合成拿到最终掩码后合成新背景def composite(image_path, mask, bg_color(255, 255, 255)): image cv2.imread(image_path) mask_3c cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR) / 255.0 bg np.full_like(image, bg_color, dtypenp.uint8) result (image * mask_3c bg * (1 - mask_3c)).astype(np.uint8) return result如果要换的是图片背景而非纯色把bg换成 resize 到同尺寸的背景图即可。注意掩码边缘可以做 12 像素的高斯模糊合成后过渡更自然不会有锯齿感。6. 避坑与排查人物抠图训练中最容易翻车的 5 个地方现象一训练 loss 正常下降但推理时掩码全白或全黑。原因训练时用了activationNone推理时忘了加 sigmoid直接把 logits 当概率用。logits 范围是负无穷到正无穷阈值化后自然全白或全黑。 解决推理代码里确认有torch.sigmoid(logits)这一步或者加载模型时设activationsigmoid但训练时必须为 None。现象二验证集 IoU 很高但实际图片抠出来边缘毛糙、有大量小碎块。原因验证集和训练集来自同一分布且验证时只做了 resize没有覆盖真实场景的多样性。另外模型输出的概率图本身就有噪声。 解决后处理加形态学开闭运算和最大连通域保留。如果碎块仍然多检查训练数据里是否有标注质量差的样本把 IoU 低于 0.5 的标注挑出来修正或丢弃。现象三训练到 10 轮左右 loss 突然变成 NaN。原因学习率过大导致梯度爆炸或者 Dice Loss 在预测全零时出现除零。另外混合精度训练AMP下 BCE 有时会溢出。 解决先把学习率降到 1e-4加梯度裁剪torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)。Dice Loss 里加smooth1.0参数避免除零。如果用了 AMP把 loss 计算部分强制转成 float32。现象四换编码器后显存不够batch_size 降到 2 还是 OOM。原因不同编码器的参数量和中间特征图大小差异很大。efficientnet-b5 比 resnet34 的显存占用高好几倍。 解决先降输入分辨率320→256再考虑换更小的编码器。另外检查是否在验证时忘了加torch.no_grad()这会导致验证阶段也构建计算图显存翻倍。现象五人物边缘的头发丝区域总是被切成背景。原因头发丝区域像素占比极小Dice Loss 对这种细结构不敏感。另外下采样过程中细节信息丢失严重。 解决换用decoder_attention_typescse增强细节关注损失函数里加边界加权对掩码边缘像素给更高权重或者换mit_b0这类 Transformer 编码器它对长距离依赖和细结构的分割效果更好。如果还不行考虑用 matting 方法做后处理精修。7. 把抠图精度再推一步从分割到 matting 的衔接思路分割模型输出的掩码是硬边界的阈值化之后只有 0 和 255 两个值。但真实的人物边缘——尤其是头发、半透明衣物、运动模糊区域——需要 alpha 渐变才能自然合成。这就是 matting 要解决的问题。一个实用的衔接方案是用 smp 分割模型生成粗掩码然后在掩码边缘的过渡带膨胀掩码减去腐蚀掩码得到的区域内用 KNN matting 或 Closed-form matting 做精修。这样既利用了分割模型的语义理解能力又用 matting 补上了边缘细节。def trimap_from_mask(mask, erode_size5, dilate_size10): # 生成三分图确定前景255确定背景0未知区域128 kernel_e cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (erode_size, erode_size)) kernel_d cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (dilate_size, dilate_size)) fg cv2.erode(mask, kernel_e, iterations1) # 腐蚀得到确定前景 bg cv2.dilate(mask, kernel_d, iterations1) # 膨胀得到确定背景 trimap np.full_like(mask, 128) # 默认未知 trimap[bg 0] 0 # 确定背景 trimap[fg 255] 255 # 确定前景 return trimap拿到 trimap 后可以用pymatting库做 alpha 估计from pymatting import estimate_alpha_cf def matting_refine(image, trimap): # image 归一化到 0~1 img_float image.astype(np.float64) / 255.0 alpha estimate_alpha_cf(img_float, trimap) return (alpha * 255).astype(np.uint8)estimate_alpha_cf是 Closed-form matting 的实现对头发丝区域效果明显。代价是速度慢一张 1080P 图大概要几秒。如果做批量处理可以只对 trimap 的未知区域做计算确定区域直接赋值能省不少时间。验证 matting 效果的方法把 alpha 掩码和原图做合成换一个纯色背景放大看边缘过渡是否自然。如果边缘有白边或黑边说明 trimap 的腐蚀/膨胀参数需要调——白边通常是确定前景范围太小黑边是确定背景范围太小。我自己做人物抠图项目时习惯在训练完分割模型后先跑一批测试图把 IoU 低于 0.85 的样本单独拎出来看。十有八九是标注问题或者极端光照条件而不是模型本身不行。把这几张图的标注修一修重新训一轮比调参折腾半天管用得多。希望帮到你。本文还有配套的精品资源点击获取
