拓冰网站定制 
示例工程【免费下载链接】deep-learning-for-image-processingdeep learning for image processing including classification and object-detection etc.项目地址https://gitcode.com/gh_mirrors/de/deep-learning-for-image-processing点击查看免费下载本篇技术指南以 deep-learning-for-image-processing 仓库中 pytorch_segmentation/u2net 模块为核心系统讲解 U2-NetGoing Deeper with Nested U-Structure for Salient Object Detection在显著性目标检测SOD任务上的完整落地流程。文章将覆盖网络结构RSU / RSU4F 源码级拆解、DUTS 数据集准备、单机与多 GPU 训练指令、官方权重加载与指标复现、post_norm 后处理差异等关键细节。读完本文你将能够在本仓库基础上独立完成 U2-Net 从数据准备、训练、验证到推理预测的完整流程并理解其在 DUTS-TE 上 MAE / maxF1 指标背后的实现原理。一、项目背景与定位U2-Net论文名Going Deeper with Nested U-Structure for Salient Object Detection是显著性目标检测Salient Object Detection, SOD领域的经典网络。与常规的编码器-解码器分割网络不同U2-Net 的核心创新在于嵌套 U 型结构在每个编码/解码阶段内部再嵌套一个 U 型子结构使网络在保持较少参数的同时能够捕捉多尺度上下文信息从而在复杂背景下也能准确分离显著目标。本仓库中的 pytorch_segmentation/u2net 目录主要移植自官方开源实现并针对训练与验证流程做了本地化改造自定义数据集读取、单/多 GPU 训练脚本、独立的验证与预测脚本、官方权重转换脚本等。需要特别注意的是该项目面向的是显著性目标检测SOD任务而不是语义分割Semantic Segmentation任务两者在标签形式二值 Mask、损失函数二元交叉熵、评价指标MAE、maxF1上均有本质区别使用时应与仓库中其他分割模块如 unet、deeplab_v3区分开。二、环境配置与依赖README 中给出的推荐环境如下项推荐配置Python3.6 / 3.7 / 3.8PyTorch1.10操作系统Ubuntu 或 CentOSWindows 暂不支持多 GPU 训练硬件建议使用 GPU 训练更精确的依赖版本以 pytorch_segmentation/u2net/requirements.txt 为准numpy1.22.0 torch1.13.1 torchvision0.11.1 opencv_python4.5.4.60从源码依赖看训练/验证核心涉及 torch模型与分布式训练、torchvision数据变换与transforms基类、opencvmy_dataset.py 中使用cv2.imread读取图片与 Mask、numpypredict.py 中预测结果后处理与 matplotlib预测可视化。环境提示README 标注环境为 Pytorch1.10而 requirements.txt 固定 torch1.13.1实际安装时应以 requirements.txt 为基准如使用更高版本 PyTorchtorch.cuda.amp相关 API 基本兼容但建议保持与仓库验证过的一致版本。三、仓库文件结构一览pytorch_segmentation/u2net 目录结构如下├── src/ # 搭建网络相关代码model.py 定义 U2-Net ├── train_utils/ # 训练以及验证相关代码train_and_eval.py、distributed_utils.py ├── my_dataset.py # 自定义数据集读取相关代码DUTS 数据集 ├── predict.py # 简易的预测代码 ├── train.py # 单GPU或CPU训练代码 ├── train_multi_GPU.py # 多GPU并行训练代码 ├── validation.py # 单独验证模型相关代码 ├── transforms.py # 数据预处理相关代码 ├── convert_weight.py # 官方权重 - 本仓库权重格式转换脚本 ├── requirements.txt # 项目依赖 └── results*.txt # 训练过程指标记录如 results20220723-123632.txt各文件职责与 README 描述一致网络搭建集中在src/训练与验证逻辑在train_utils/数据读取、预处理、推理分别由my_dataset.py、transforms.py、predict.py承担。额外值得一提的文件是 convert_weight.py它用于把官方权重键名形如stage1.rebnconvin.conv.weight转换为本仓库模型键名形如encode_modules.0.conv_in.conv.weight可加载的格式并提供了新旧权重逐层映射encode/decode 各阶段、side 输出、out_conv与输出一致性断言是从官方模型迁移权重的关键工具。四、DUTS 数据集准备U2-Net 论文使用 DUTS 数据集训练与评测本仓库同样以 DUTS 为标准数据源。4.1 数据集获取与目录结构DUTS 数据集由 DUTS-TR训练集与 DUTS-TE测试/验证集两部分组成可在 DUTS 官方数据发布页获取README 中同时提供了作者维护的镜像下载链接。解压后目录结构如下├── DUTS-TR │ ├── DUTS-TR-Image: 该文件夹存放所有训练集的图片 │ └── DUTS-TR-Mask: 该文件夹存放对应训练图片的GT标签Mask蒙板形式 │ └── DUTS-TE ├── DUTS-TE-Image: 该文件夹存放所有测试验证集的图片 └── DUTS-TE-Mask: 该文件夹存放对应测试验证图片的GT标签Mask蒙板形式注意训练或者验证过程中将--data-path指向DUTS-TR 所在根目录即同时包含 DUTS-TR 与 DUTS-TE 两个子文件夹的目录脚本会根据 train 标志自动拼接子路径。4.2 源码级读取逻辑my_dataset.py 中DUTSDataset的目录约定与上述结构完全一致训练模式trainTrue图片根目录为DUTS-TR/DUTS-TR-ImageMask 根目录为DUTS-TR/DUTS-TR-Mask验证模式trainFalse图片根目录为DUTS-TE/DUTS-TE-ImageMask 根目录为DUTS-TE/DUTS-TE-Mask。读取时图片按.jpg后缀收集Mask 按.png后缀收集并通过图片名.replace(.jpg, .png)建立一一对应关系若某张图片缺少对应 Mask 会直接assert报错——这保证了训练样本的图文对齐。图像读取采用 OpenCV图片以IMREAD_COLOR读入后经cv2.cvtColor由 BGR 转为 RGBMask 以IMREAD_GRAYSCALE读取为单通道灰度图。由于训练时随机裁剪会造成同 batch 内尺寸不一致my_dataset.py 提供了cat_list静态方法取 batch 内最大尺寸创建填充张量将每张图/每个 Mask 拷贝到左上角其余区域以 0 填充作为collate_fn使用。五、U2-Net 网络结构源码解析U2-Net 网络结构如上图所示整体呈大 U 套小 U的嵌套形态。本仓库在 src/model.py 中完整实现了该结构可以对照源码理解每一层的真实含义。5.1 基础模块ConvBNReLU / DownConvBNReLU / UpConvBNReLUclass ConvBNReLU(nn.Module): def __init__(self, in_ch: int, out_ch: int, kernel_size: int 3, dilation: int 1): super().__init__() padding kernel_size // 2 if dilation 1 else dilation self.conv nn.Conv2d(in_ch, out_ch, kernel_size, paddingpadding, dilationdilation, biasFalse) self.bn nn.BatchNorm2d(out_ch) self.relu nn.ReLU(inplaceTrue) def forward(self, x: torch.Tensor) - torch.Tensor: return self.relu(self.bn(self.conv(x)))卷积默认kernel_size3支持dilation空洞卷积参数padding 规则为dilation1时padding kernel_size // 2否则padding dilation保证输出尺寸不缩。卷积 bias 默认为 False因为卷积后紧跟 BN 层bias 不起实际作用。这一细节与官方权重加载有关详见第六节。DownConvBNReLU与UpConvBNReLU继承自ConvBNReLU分别代表下采样可选F.max_pool2d(kernel2, stride2, ceil_modeTrue)与上采样拼接F.interpolate双线性插值到目标尺寸后与另一分支torch.cat拼接再卷积。注意这里的下采样/上采样是嵌套子 U 结构内部的采样外层编解码之间的下采样F.max_pool2d与上采样F.interpolate在U2Net.forward中另行处理。5.2 核心构建块RSU 与 RSU4FRSUReSidual U-block是 U2-Net 的基本构建单元其构造参数为(height, in_ch, mid_ch, out_ch)class RSU(nn.Module): def __init__(self, height: int, in_ch: int, mid_ch: int, out_ch: int): super().__init__() assert height 2 self.conv_in ConvBNReLU(in_ch, out_ch) encode_list [DownConvBNReLU(out_ch, mid_ch, flagFalse)] decode_list [UpConvBNReLU(mid_ch * 2, mid_ch, flagFalse)] for i in range(height - 2): encode_list.append(DownConvBNReLU(mid_ch, mid_ch)) decode_list.append(UpConvBNReLU(mid_ch * 2, mid_ch if i height - 3 else out_ch)) encode_list.append(ConvBNReLU(mid_ch, mid_ch, dilation2)) self.encode_modules nn.ModuleList(encode_list) self.decode_modules nn.ModuleList(decode_list) def forward(self, x: torch.Tensor) - torch.Tensor: x_in self.conv_in(x) x x_in encode_outputs [] for m in self.encode_modules: x m(x) encode_outputs.append(x) x encode_outputs.pop() for m in self.decode_modules: x2 encode_outputs.pop() x m(x, x2) return x x_in结构要点conv_in先做一次 3×3 卷积升到out_ch编码路径由height-1个下采样块 1 个空洞卷积dilation2块构成深度由height决定如 RSU-7 内部有 6 层下采样 1 层空洞卷积解码路径对称将编码层输出与解码层特征逐级拼接UpConvBNReLU中torch.cat([x1, x2], dim1)因此输入通道为mid_ch*2最终return x x_in是残差连接Residual U-block 的由来。RSU4F是 RSU 的变体4F 表示仅使用 4 层空洞卷积不使用池化下采样依靠 dilation 分别为 2、4、8 的空洞卷积扩大感受野适合放在网络最深处stage5/6以获取大范围上下文同时保持特征图分辨率class RSU4F(nn.Module): def __init__(self, in_ch: int, mid_ch: int, out_ch: int): super().__init__() self.conv_in ConvBNReLU(in_ch, out_ch) self.encode_modules nn.ModuleList([ConvBNReLU(out_ch, mid_ch), ConvBNReLU(mid_ch, mid_ch, dilation2), ConvBNReLU(mid_ch, mid_ch, dilation4), ConvBNReLU(mid_ch, mid_ch, dilation8)]) self.decode_modules nn.ModuleList([ConvBNReLU(mid_ch * 2, mid_ch, dilation4), ConvBNReLU(mid_ch * 2, mid_ch, dilation2), ConvBNReLU(mid_ch * 2, out_ch)])与 RSU 一样RSU4F 也通过x x_in保留残差连接。5.3 整体网络U2Net 与 u2net_full / u2net_lite 配置U2Net类通过配置文件cfg描述所有编解码阶段每项配置长度为 6[height, in_ch, mid_ch, out_ch, 是否RSU4F, 是否输出side]。u2net_full()的完整配置来自 src/model.py# height, in_ch, mid_ch, out_ch, RSU4F, side encode: [[7, 3, 32, 64, False, False], # En1 [6, 64, 32, 128, False, False], # En2 [5, 128, 64, 256, False, False], # En3 [4, 256, 128, 512, False, False], # En4 [4, 512, 256, 512, True, False], # En5 [4, 512, 256, 512, True, True]], # En6 decode: [[4, 1024, 256, 512, True, True], # De5 [4, 1024, 128, 256, False, True], # De4 [5, 512, 64, 128, False, True], # De3 [6, 256, 32, 64, False, True], # De2 [7, 128, 16, 64, False, True]] # De1u2net_lite()轻量版则统一将各阶段通道压缩到 64 以内mid_ch统一为 16显著减少参数量配置同样位于 src/model.py。U2Net.forward的前向流程src/model.py清晰展示了大 U的完整链路编码阶段依次经过 6 个 encode 模块除最后一个模块外每个模块输出后接F.max_pool2d(kernel2, stride2, ceil_modeTrue)下采样解码阶段从最深层特征开始每次F.interpolate双线性上采样到对应编码层尺寸与编码输出torch.concat拼接后送入 decode 模块并将结果按序插入decode_outputs前端侧输出side output配置中sideTrue的 6 个阶段En6、De5、De4、De3、De2、De1分别接一个 3×3 卷积将特征图映射到单通道再插值回输入尺寸[h, w]融合输出6 个侧输出在通道维torch.concat后经 1×1 卷积out_conv融合为最终单通道显著图。训练与推理分支不同self.training为 True 时返回[x] side_outputs即融合输出 6 个侧输出共 7 个图用于深度监督且不经过 sigmoid配合binary_cross_entropy_with_logits以保证 AMP 混合精度安全推理时仅返回torch.sigmoid(x)融合输出。5.4 损失函数与深度监督src/model.py 中 U2-Net 在训练时输出融合图与 6 个侧输出共 7 个预测图。对应损失函数定义在 train_utils/train_and_eval.py 的criteriondef criterion(inputs, target): losses [F.binary_cross_entropy_with_logits(inputs[i], target) for i in range(len(inputs))] total_loss sum(losses) return total_loss即对融合输出与每一个侧输出分别计算二元交叉熵BCE损失并求和这就是 U2-Net 深度监督deep supervision的训练方式——每个侧输出都会引导网络在对应尺度上学习显著目标结构。六、官方权重加载bias 细节与 post_norm 差异6.1 权重来源与转换README 提供了从官方转换得到的两份权重u2net_full.pth完整版与u2net_lite.pth轻量版。如需自行从官方权重转换仓库提供了 convert_weight.py它将官方权重中stageX.rebnconvX[.dX]形式的键名按 encode/decode 阶段映射为本仓库encode_modules.{i}.encode_modules.{j}/decode_modules.{i}.decode_modules.{j}等形式的键名side 输出与out_conv也有对应映射并支持strictTrue严格加载校验。6.2 关键细节ConvBNReLU 的 bias 必须为 TrueREADME 明确指出如果要载入官方提供的权重需要将 src/model.py 中ConvBNReLU类里卷积的 bias 设置成 True因为官方代码里没有进行显式设置PyTorchConv2d的 bias 默认为 True即官方权重中存在conv.bias参数。而由于卷积后紧跟 BNbias 实际起不到作用本仓库默认将 bias 设置为 False因此直接加载官方权重会因键不匹配而失败。两种处理方式修改源码 bias 为 True或使用 convert_weight.py 转换均可解决推荐后者无需改动模型代码。6.3 post_norm 后处理与指标差异README 用validation.py验证u2net_full在 DUTS-TE 上的结果MAE: 0.044 maxF1: 0.868并特别说明这里的 maxF1 与原论文结果存在差异差异主要来自 post_norm。原仓库会对预测结果做 post_norm而本仓库在 train_utils/train_and_eval.py 的evaluate中将 post_norm 注释移除了。如果需要恢复该后处理可自行添加其中output为验证时网络预测的输出ma torch.max(output) mi torch.min(output) output (output - mi) / (ma - mi)即把预测图线性归一化到 [0, 1] 区间。README 说明加上 post_norm 后 maxF1 约为0.872。6.4 指标实现源码验证使用的两个指标定义在 train_utils/distributed_utils.pyMeanAbsoluteError将预测插值到 GT 尺寸后计算sum(|pred - gt|) / (h*w)即逐像素绝对误差均值MAE 越小越好F1Score参考显著性检测社区实现按预测概率直方图从大到小累积计算 Precision / Recall 曲线并以加权调和均值F1 (10.3)*P*R / (0.3*PR)β²0.3在曲线上取最大值作为 maxF1。两者均要求验证 batch_size 必须为 1MeanAbsoluteError.update与F1Score.update中有assert batch_size 1这也是 validation.py 与两个训练脚本中验证 DataLoaderbatch_size1的原因。七、训练实战单 GPU 与多 GPU7.1 数据预处理训练/验证预处理定义在 transforms.py 中均对图片与 Mask 同步操作训练SODPresetTrain见 train.pyToTensor→Resize([320,320], resize_maskTrue)图片与 Mask 同步缩放到 320×320→RandomCrop(288)随机裁剪 288×288→RandomHorizontalFlip(0.5)→Normalize(mean(0.485,0.456,0.406), std(0.229,0.224,0.225))验证SODPresetEvalToTensor→Resize([320,320], resize_maskFalse)Mask 不缩放保持原 GT 分辨率便于指标计算→Normalize。其中RandomCrop在图像小于裁剪尺寸时会先用 0 填充pad_if_smaller保证任意输入尺寸可裁剪。7.2 单 GPU / CPU 训练train.py直接使用 train.py 即可脚本会自动根据torch.cuda.is_available()决定设备--device cuda但无 GPU 时自动回退 CPUpython train.py --data-path /path/to/DUTS_root主要命令行参数train.py参数默认值说明--data-path./DUTS 根目录包含 DUTS-TR 与 DUTS-TE--devicecuda训练设备-b / --batch-size16batch size--epochs360总训练轮数--lr0.001初始学习率--wd / --weight-decay1e-4权重衰减--eval-interval10每间隔多少个 epoch 验证一次--print-freq50打印频率--resume从 checkpoint 恢复训练--start-epoch0起始 epoch--ampFalseflag开启 torch.cuda.amp 混合精度训练训练相关实现细节优化器为AdamW配合get_params_groupstrain_utils/train_and_eval.py将 BN 参数与 bias 从权重衰减中剔除学习率采用warmup cosine 衰减调度create_lr_scheduler默认 warmup 2 个 epoch倍率因子从 1e-3 升至 1再沿余弦曲线降至 1e-6开启--amp时使用torch.cuda.amp.GradScaler与autocast损失计算采用binary_cross_entropy_with_logits不经过 sigmoid正是为了保证混合精度下的数值安全训练日志写入results{时间戳}.txt每--eval-interval个 epoch 验证一次并记录train_loss / lr / MAE / maxF1权重保存到save_weights/目录仅保留最近 10 个 epoch 的权重文件并在 MAE 不升、maxF1 不降时额外保存model_best.pth。7.3 多 GPU 训练train_multi_GPU.py使用 PyTorch 官方torchrun启动torchrun --nproc_per_node4 train_multi_GPU.py --lr 0.004 --amp--nproc_per_node为使用 GPU 数量若只想使用部分 GPU 设备可在指令前加上CUDA_VISIBLE_DEVICES环境变量指定设备编号例如只使用物理设备第 1 块与第 4 块CUDA_VISIBLE_DEVICES0,3 torchrun --nproc_per_node2 train_multi_GPU.pytrain_multi_GPU.py 与单机版的主要差异通过init_distributed_modetrain_utils/distributed_utils.py读取RANK / WORLD_SIZE / LOCAL_RANK环境变量完成分布式初始化使用DistributedSampler划分数据DistributedDataParallel包装模型可选--sync-bn开启同步 BN多卡间同步训练会变慢参数略有扩展-j / --workers默认 4、--output-dir默认./multi_train、--world-size、--dist-url默认env://、--test-only仅测试不训练等其中--batch-size表示每块 GPU 上的 batch size全局 batch size 为NGPU × batch_size日志写入与权重保存仅在 rank 0 主进程执行save_on_master。注意README 注明 Windows 暂不支持多 GPU 训练torchrun与nccl后端依赖 Linux 环境多 GPU 训练建议在 Ubuntu / CentOS 上进行。八、训练记录与结果对比README 记录了u2net_full在本仓库配置下的训练结果。训练指令4 卡、学习率 0.004、开启 AMPtorchrun --nproc_per_node4 train_multi_GPU.py --lr 0.004 --amp最终在 DUTS-TE 上的验证结果MAE: 0.047 maxF1: 0.859与官方权重验证结果MAE 0.044 / maxF1 0.868相比略有差距符合正常复现水平差异主要来自 post_norm 后处理与训练细节。训练过程详情可参考仓库中的 results20220723-123632.txt从该记录可以看到典型的收敛曲线第 0 个 epochtrain_loss: 2.7385 lr: 0.002002 MAE: 0.465 maxF1: 0.464warmup 初期第 100 个 epochMAE: 0.055 maxF1: 0.847第 190 个 epoch 附近进入平台期MAE: 0.047 maxF1: 0.860最终第 359 个 epochtrain_loss: 0.1518 lr: 0.000000 MAE: 0.047 maxF1: 0.859。整个训练过程学习率从 warmup 后约 0.004 沿余弦曲线平滑衰减到 0MAE 从 0.465 稳步降至 0.047 左右maxF1 从 0.464 提升至约 0.859验证集指标在后期趋于稳定——这是判断训练收敛与是否需要早停的直接依据。九、验证与推理9.1 单独验证validation.py使用 validation.py 对指定权重在 DUTS-TE 上单独验证python validation.py --data-path /path/to/DUTS_root --weights ./u2net_full.pth参数说明--data-path为 DUTS 根目录--weights指定权重文件路径--device默认cuda:0。脚本构建SODPresetEval([320, 320])预处理后以 batch_size1 逐张送入网络由evaluate计算 MAE 与 maxF1 并打印。权重加载兼容两种格式直接 load_state_dict或从训练保存的含model键的 checkpoint 中提取validation.py。9.2 单图预测predict.pypredict.py 提供了最简推理流程weights_path ./u2net_full.pth img_path ./test.png threshold 0.5执行python predict.py前将weights_path指向权重、img_path指向测试图片即可。其推理流程要点图片 BGR→RGB 后依次做ToTensor → Resize(320) → NormalizeImageNet 均值/方差加载u2net_full()权重并model.eval()先用全零init_img预热模型触发 CUDA kernel 初始化再对真实图片推理并统计耗时推理输出torch.sigmoid后的单通道显著图缩放到原图尺寸按threshold0.5二值化生成 Mask用显著区域origin_img * pred_mask提取前景并保存pred_result.png同时用 matplotlib 显示。9.3 ONNX 导出src/model.py 末尾还附带convert_onnx函数可将 RSU / RSU4F 或完整 U2-Net 以opset_version11导出为 ONNX输入尺寸 288×288方便在部署侧使用该仓库另有独立的模型部署章节可供参考。十、总结与常见注意事项综合 README 与源码使用本仓库训练/使用 U2-Net 时需要重点注意以下几点任务定位U2-Net 属于显著性目标检测SOD标签为二值 Mask评估指标是 MAE 与 maxF1与语义分割任务不可混用数据集结构--data-path必须指向包含DUTS-TR与DUTS-TE的根目录且训练图片为.jpg、Mask 为同名.png官方权重加载官方权重包含conv.bias键而本仓库ConvBNReLU默认biasFalse可用 convert_weight.py 转换权重或按 README 说明将 bias 改为 True指标差异来源本仓库验证时移除了 post_norm(output - mi) / (ma - mi)导致 maxF1 相比原论文约低 0.0040.868 vs 0.872如需复现原论文数值可自行恢复该后处理多 GPU 训练使用torchrun --nproc_per_nodeN train_multi_GPU.py配合CUDA_VISIBLE_DEVICES指定 GPU仅支持 LinuxWindows 不支持多 GPU验证 batch_sizeMAE / F1 指标实现要求验证时 batch_size 必须为 1脚本已固定无需修改深度监督损失训练时融合输出与 6 个侧输出各计算一次 BCE 并求和这是 U2-Net 多尺度学习能力的关键也是 src/model.py 中self.training分支存在的意义。至此从网络结构、数据集、训练、验证到推理的完整链路已全部打通。若需进一步深入理解 U2-Net 代码细节与训练策略可直接阅读 src/model.py、train_utils/train_and_eval.py 与 train_utils/distributed_utils.py 的完整实现并参照仓库内results*.txt复现训练过程。赞分享示例工程【免费下载链接】deep-learning-for-image-processingdeep learning for image processing including classification and object-detection etc.项目地址https://gitcode.com/gh_mirrors/de/deep-learning-for-image-processing点击查看免费下载相关推荐U-Net 图像分割实战基于 deep-learning-for-image-processing 仓库的 DRIVE 视网膜血管分割与 PyTorch 训练部署指南U Net 图像分割实战基于 deep learning for image processing 仓库的 DRIVE 视网膜血管分割与 PyTorch 训练示例工程Colorify.js扩展开发指南如何自定义插件和增强颜色提取功能Colorify.js扩展开发指南如何自定义插件和增强颜色提取功能 Colorify.js是一个简单、可定制、轻量级的JavaScript颜色提取库能够从图Mask R-CNN 实例分割实战基于 deep-learning-for-image-processing 仓库的 COCO / Pascal VOC 训练、验证与部署指南Mask R CNN 实例分割实战基于 deep learning for image processing 仓库的 COCO / Pascal VOC 训练示例工程上一篇CloddsBot多语言设置10种语言i18n切换完整教程下一篇Sketch MeaXure重新定义设计标注工作流的现代解决方案创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 deep-learning-for-image-processing 仓库 U2-Net 实战：基于嵌套 U 型结构的显著性目标检测（SOD）训练与推理全解析
示例工程【免费下载链接】deep-learning-for-image-processingdeep learning for image processing including classification and object-detection etc.项目地址https://gitcode.com/gh_mirrors/de/deep-learning-for-image-processing点击查看免费下载本篇技术指南以 deep-learning-for-image-processing 仓库中 pytorch_segmentation/u2net 模块为核心系统讲解 U2-NetGoing Deeper with Nested U-Structure for Salient Object Detection在显著性目标检测SOD任务上的完整落地流程。文章将覆盖网络结构RSU / RSU4F 源码级拆解、DUTS 数据集准备、单机与多 GPU 训练指令、官方权重加载与指标复现、post_norm 后处理差异等关键细节。读完本文你将能够在本仓库基础上独立完成 U2-Net 从数据准备、训练、验证到推理预测的完整流程并理解其在 DUTS-TE 上 MAE / maxF1 指标背后的实现原理。一、项目背景与定位U2-Net论文名Going Deeper with Nested U-Structure for Salient Object Detection是显著性目标检测Salient Object Detection, SOD领域的经典网络。与常规的编码器-解码器分割网络不同U2-Net 的核心创新在于嵌套 U 型结构在每个编码/解码阶段内部再嵌套一个 U 型子结构使网络在保持较少参数的同时能够捕捉多尺度上下文信息从而在复杂背景下也能准确分离显著目标。本仓库中的 pytorch_segmentation/u2net 目录主要移植自官方开源实现并针对训练与验证流程做了本地化改造自定义数据集读取、单/多 GPU 训练脚本、独立的验证与预测脚本、官方权重转换脚本等。需要特别注意的是该项目面向的是显著性目标检测SOD任务而不是语义分割Semantic Segmentation任务两者在标签形式二值 Mask、损失函数二元交叉熵、评价指标MAE、maxF1上均有本质区别使用时应与仓库中其他分割模块如 unet、deeplab_v3区分开。二、环境配置与依赖README 中给出的推荐环境如下项推荐配置Python3.6 / 3.7 / 3.8PyTorch1.10操作系统Ubuntu 或 CentOSWindows 暂不支持多 GPU 训练硬件建议使用 GPU 训练更精确的依赖版本以 pytorch_segmentation/u2net/requirements.txt 为准numpy1.22.0 torch1.13.1 torchvision0.11.1 opencv_python4.5.4.60从源码依赖看训练/验证核心涉及 torch模型与分布式训练、torchvision数据变换与transforms基类、opencvmy_dataset.py 中使用cv2.imread读取图片与 Mask、numpypredict.py 中预测结果后处理与 matplotlib预测可视化。环境提示README 标注环境为 Pytorch1.10而 requirements.txt 固定 torch1.13.1实际安装时应以 requirements.txt 为基准如使用更高版本 PyTorchtorch.cuda.amp相关 API 基本兼容但建议保持与仓库验证过的一致版本。三、仓库文件结构一览pytorch_segmentation/u2net 目录结构如下├── src/ # 搭建网络相关代码model.py 定义 U2-Net ├── train_utils/ # 训练以及验证相关代码train_and_eval.py、distributed_utils.py ├── my_dataset.py # 自定义数据集读取相关代码DUTS 数据集 ├── predict.py # 简易的预测代码 ├── train.py # 单GPU或CPU训练代码 ├── train_multi_GPU.py # 多GPU并行训练代码 ├── validation.py # 单独验证模型相关代码 ├── transforms.py # 数据预处理相关代码 ├── convert_weight.py # 官方权重 - 本仓库权重格式转换脚本 ├── requirements.txt # 项目依赖 └── results*.txt # 训练过程指标记录如 results20220723-123632.txt各文件职责与 README 描述一致网络搭建集中在src/训练与验证逻辑在train_utils/数据读取、预处理、推理分别由my_dataset.py、transforms.py、predict.py承担。额外值得一提的文件是 convert_weight.py它用于把官方权重键名形如stage1.rebnconvin.conv.weight转换为本仓库模型键名形如encode_modules.0.conv_in.conv.weight可加载的格式并提供了新旧权重逐层映射encode/decode 各阶段、side 输出、out_conv与输出一致性断言是从官方模型迁移权重的关键工具。四、DUTS 数据集准备U2-Net 论文使用 DUTS 数据集训练与评测本仓库同样以 DUTS 为标准数据源。4.1 数据集获取与目录结构DUTS 数据集由 DUTS-TR训练集与 DUTS-TE测试/验证集两部分组成可在 DUTS 官方数据发布页获取README 中同时提供了作者维护的镜像下载链接。解压后目录结构如下├── DUTS-TR │ ├── DUTS-TR-Image: 该文件夹存放所有训练集的图片 │ └── DUTS-TR-Mask: 该文件夹存放对应训练图片的GT标签Mask蒙板形式 │ └── DUTS-TE ├── DUTS-TE-Image: 该文件夹存放所有测试验证集的图片 └── DUTS-TE-Mask: 该文件夹存放对应测试验证图片的GT标签Mask蒙板形式注意训练或者验证过程中将--data-path指向DUTS-TR 所在根目录即同时包含 DUTS-TR 与 DUTS-TE 两个子文件夹的目录脚本会根据 train 标志自动拼接子路径。4.2 源码级读取逻辑my_dataset.py 中DUTSDataset的目录约定与上述结构完全一致训练模式trainTrue图片根目录为DUTS-TR/DUTS-TR-ImageMask 根目录为DUTS-TR/DUTS-TR-Mask验证模式trainFalse图片根目录为DUTS-TE/DUTS-TE-ImageMask 根目录为DUTS-TE/DUTS-TE-Mask。读取时图片按.jpg后缀收集Mask 按.png后缀收集并通过图片名.replace(.jpg, .png)建立一一对应关系若某张图片缺少对应 Mask 会直接assert报错——这保证了训练样本的图文对齐。图像读取采用 OpenCV图片以IMREAD_COLOR读入后经cv2.cvtColor由 BGR 转为 RGBMask 以IMREAD_GRAYSCALE读取为单通道灰度图。由于训练时随机裁剪会造成同 batch 内尺寸不一致my_dataset.py 提供了cat_list静态方法取 batch 内最大尺寸创建填充张量将每张图/每个 Mask 拷贝到左上角其余区域以 0 填充作为collate_fn使用。五、U2-Net 网络结构源码解析U2-Net 网络结构如上图所示整体呈大 U 套小 U的嵌套形态。本仓库在 src/model.py 中完整实现了该结构可以对照源码理解每一层的真实含义。5.1 基础模块ConvBNReLU / DownConvBNReLU / UpConvBNReLUclass ConvBNReLU(nn.Module): def __init__(self, in_ch: int, out_ch: int, kernel_size: int 3, dilation: int 1): super().__init__() padding kernel_size // 2 if dilation 1 else dilation self.conv nn.Conv2d(in_ch, out_ch, kernel_size, paddingpadding, dilationdilation, biasFalse) self.bn nn.BatchNorm2d(out_ch) self.relu nn.ReLU(inplaceTrue) def forward(self, x: torch.Tensor) - torch.Tensor: return self.relu(self.bn(self.conv(x)))卷积默认kernel_size3支持dilation空洞卷积参数padding 规则为dilation1时padding kernel_size // 2否则padding dilation保证输出尺寸不缩。卷积 bias 默认为 False因为卷积后紧跟 BN 层bias 不起实际作用。这一细节与官方权重加载有关详见第六节。DownConvBNReLU与UpConvBNReLU继承自ConvBNReLU分别代表下采样可选F.max_pool2d(kernel2, stride2, ceil_modeTrue)与上采样拼接F.interpolate双线性插值到目标尺寸后与另一分支torch.cat拼接再卷积。注意这里的下采样/上采样是嵌套子 U 结构内部的采样外层编解码之间的下采样F.max_pool2d与上采样F.interpolate在U2Net.forward中另行处理。5.2 核心构建块RSU 与 RSU4FRSUReSidual U-block是 U2-Net 的基本构建单元其构造参数为(height, in_ch, mid_ch, out_ch)class RSU(nn.Module): def __init__(self, height: int, in_ch: int, mid_ch: int, out_ch: int): super().__init__() assert height 2 self.conv_in ConvBNReLU(in_ch, out_ch) encode_list [DownConvBNReLU(out_ch, mid_ch, flagFalse)] decode_list [UpConvBNReLU(mid_ch * 2, mid_ch, flagFalse)] for i in range(height - 2): encode_list.append(DownConvBNReLU(mid_ch, mid_ch)) decode_list.append(UpConvBNReLU(mid_ch * 2, mid_ch if i height - 3 else out_ch)) encode_list.append(ConvBNReLU(mid_ch, mid_ch, dilation2)) self.encode_modules nn.ModuleList(encode_list) self.decode_modules nn.ModuleList(decode_list) def forward(self, x: torch.Tensor) - torch.Tensor: x_in self.conv_in(x) x x_in encode_outputs [] for m in self.encode_modules: x m(x) encode_outputs.append(x) x encode_outputs.pop() for m in self.decode_modules: x2 encode_outputs.pop() x m(x, x2) return x x_in结构要点conv_in先做一次 3×3 卷积升到out_ch编码路径由height-1个下采样块 1 个空洞卷积dilation2块构成深度由height决定如 RSU-7 内部有 6 层下采样 1 层空洞卷积解码路径对称将编码层输出与解码层特征逐级拼接UpConvBNReLU中torch.cat([x1, x2], dim1)因此输入通道为mid_ch*2最终return x x_in是残差连接Residual U-block 的由来。RSU4F是 RSU 的变体4F 表示仅使用 4 层空洞卷积不使用池化下采样依靠 dilation 分别为 2、4、8 的空洞卷积扩大感受野适合放在网络最深处stage5/6以获取大范围上下文同时保持特征图分辨率class RSU4F(nn.Module): def __init__(self, in_ch: int, mid_ch: int, out_ch: int): super().__init__() self.conv_in ConvBNReLU(in_ch, out_ch) self.encode_modules nn.ModuleList([ConvBNReLU(out_ch, mid_ch), ConvBNReLU(mid_ch, mid_ch, dilation2), ConvBNReLU(mid_ch, mid_ch, dilation4), ConvBNReLU(mid_ch, mid_ch, dilation8)]) self.decode_modules nn.ModuleList([ConvBNReLU(mid_ch * 2, mid_ch, dilation4), ConvBNReLU(mid_ch * 2, mid_ch, dilation2), ConvBNReLU(mid_ch * 2, out_ch)])与 RSU 一样RSU4F 也通过x x_in保留残差连接。5.3 整体网络U2Net 与 u2net_full / u2net_lite 配置U2Net类通过配置文件cfg描述所有编解码阶段每项配置长度为 6[height, in_ch, mid_ch, out_ch, 是否RSU4F, 是否输出side]。u2net_full()的完整配置来自 src/model.py# height, in_ch, mid_ch, out_ch, RSU4F, side encode: [[7, 3, 32, 64, False, False], # En1 [6, 64, 32, 128, False, False], # En2 [5, 128, 64, 256, False, False], # En3 [4, 256, 128, 512, False, False], # En4 [4, 512, 256, 512, True, False], # En5 [4, 512, 256, 512, True, True]], # En6 decode: [[4, 1024, 256, 512, True, True], # De5 [4, 1024, 128, 256, False, True], # De4 [5, 512, 64, 128, False, True], # De3 [6, 256, 32, 64, False, True], # De2 [7, 128, 16, 64, False, True]] # De1u2net_lite()轻量版则统一将各阶段通道压缩到 64 以内mid_ch统一为 16显著减少参数量配置同样位于 src/model.py。U2Net.forward的前向流程src/model.py清晰展示了大 U的完整链路编码阶段依次经过 6 个 encode 模块除最后一个模块外每个模块输出后接F.max_pool2d(kernel2, stride2, ceil_modeTrue)下采样解码阶段从最深层特征开始每次F.interpolate双线性上采样到对应编码层尺寸与编码输出torch.concat拼接后送入 decode 模块并将结果按序插入decode_outputs前端侧输出side output配置中sideTrue的 6 个阶段En6、De5、De4、De3、De2、De1分别接一个 3×3 卷积将特征图映射到单通道再插值回输入尺寸[h, w]融合输出6 个侧输出在通道维torch.concat后经 1×1 卷积out_conv融合为最终单通道显著图。训练与推理分支不同self.training为 True 时返回[x] side_outputs即融合输出 6 个侧输出共 7 个图用于深度监督且不经过 sigmoid配合binary_cross_entropy_with_logits以保证 AMP 混合精度安全推理时仅返回torch.sigmoid(x)融合输出。5.4 损失函数与深度监督src/model.py 中 U2-Net 在训练时输出融合图与 6 个侧输出共 7 个预测图。对应损失函数定义在 train_utils/train_and_eval.py 的criteriondef criterion(inputs, target): losses [F.binary_cross_entropy_with_logits(inputs[i], target) for i in range(len(inputs))] total_loss sum(losses) return total_loss即对融合输出与每一个侧输出分别计算二元交叉熵BCE损失并求和这就是 U2-Net 深度监督deep supervision的训练方式——每个侧输出都会引导网络在对应尺度上学习显著目标结构。六、官方权重加载bias 细节与 post_norm 差异6.1 权重来源与转换README 提供了从官方转换得到的两份权重u2net_full.pth完整版与u2net_lite.pth轻量版。如需自行从官方权重转换仓库提供了 convert_weight.py它将官方权重中stageX.rebnconvX[.dX]形式的键名按 encode/decode 阶段映射为本仓库encode_modules.{i}.encode_modules.{j}/decode_modules.{i}.decode_modules.{j}等形式的键名side 输出与out_conv也有对应映射并支持strictTrue严格加载校验。6.2 关键细节ConvBNReLU 的 bias 必须为 TrueREADME 明确指出如果要载入官方提供的权重需要将 src/model.py 中ConvBNReLU类里卷积的 bias 设置成 True因为官方代码里没有进行显式设置PyTorchConv2d的 bias 默认为 True即官方权重中存在conv.bias参数。而由于卷积后紧跟 BNbias 实际起不到作用本仓库默认将 bias 设置为 False因此直接加载官方权重会因键不匹配而失败。两种处理方式修改源码 bias 为 True或使用 convert_weight.py 转换均可解决推荐后者无需改动模型代码。6.3 post_norm 后处理与指标差异README 用validation.py验证u2net_full在 DUTS-TE 上的结果MAE: 0.044 maxF1: 0.868并特别说明这里的 maxF1 与原论文结果存在差异差异主要来自 post_norm。原仓库会对预测结果做 post_norm而本仓库在 train_utils/train_and_eval.py 的evaluate中将 post_norm 注释移除了。如果需要恢复该后处理可自行添加其中output为验证时网络预测的输出ma torch.max(output) mi torch.min(output) output (output - mi) / (ma - mi)即把预测图线性归一化到 [0, 1] 区间。README 说明加上 post_norm 后 maxF1 约为0.872。6.4 指标实现源码验证使用的两个指标定义在 train_utils/distributed_utils.pyMeanAbsoluteError将预测插值到 GT 尺寸后计算sum(|pred - gt|) / (h*w)即逐像素绝对误差均值MAE 越小越好F1Score参考显著性检测社区实现按预测概率直方图从大到小累积计算 Precision / Recall 曲线并以加权调和均值F1 (10.3)*P*R / (0.3*PR)β²0.3在曲线上取最大值作为 maxF1。两者均要求验证 batch_size 必须为 1MeanAbsoluteError.update与F1Score.update中有assert batch_size 1这也是 validation.py 与两个训练脚本中验证 DataLoaderbatch_size1的原因。七、训练实战单 GPU 与多 GPU7.1 数据预处理训练/验证预处理定义在 transforms.py 中均对图片与 Mask 同步操作训练SODPresetTrain见 train.pyToTensor→Resize([320,320], resize_maskTrue)图片与 Mask 同步缩放到 320×320→RandomCrop(288)随机裁剪 288×288→RandomHorizontalFlip(0.5)→Normalize(mean(0.485,0.456,0.406), std(0.229,0.224,0.225))验证SODPresetEvalToTensor→Resize([320,320], resize_maskFalse)Mask 不缩放保持原 GT 分辨率便于指标计算→Normalize。其中RandomCrop在图像小于裁剪尺寸时会先用 0 填充pad_if_smaller保证任意输入尺寸可裁剪。7.2 单 GPU / CPU 训练train.py直接使用 train.py 即可脚本会自动根据torch.cuda.is_available()决定设备--device cuda但无 GPU 时自动回退 CPUpython train.py --data-path /path/to/DUTS_root主要命令行参数train.py参数默认值说明--data-path./DUTS 根目录包含 DUTS-TR 与 DUTS-TE--devicecuda训练设备-b / --batch-size16batch size--epochs360总训练轮数--lr0.001初始学习率--wd / --weight-decay1e-4权重衰减--eval-interval10每间隔多少个 epoch 验证一次--print-freq50打印频率--resume从 checkpoint 恢复训练--start-epoch0起始 epoch--ampFalseflag开启 torch.cuda.amp 混合精度训练训练相关实现细节优化器为AdamW配合get_params_groupstrain_utils/train_and_eval.py将 BN 参数与 bias 从权重衰减中剔除学习率采用warmup cosine 衰减调度create_lr_scheduler默认 warmup 2 个 epoch倍率因子从 1e-3 升至 1再沿余弦曲线降至 1e-6开启--amp时使用torch.cuda.amp.GradScaler与autocast损失计算采用binary_cross_entropy_with_logits不经过 sigmoid正是为了保证混合精度下的数值安全训练日志写入results{时间戳}.txt每--eval-interval个 epoch 验证一次并记录train_loss / lr / MAE / maxF1权重保存到save_weights/目录仅保留最近 10 个 epoch 的权重文件并在 MAE 不升、maxF1 不降时额外保存model_best.pth。7.3 多 GPU 训练train_multi_GPU.py使用 PyTorch 官方torchrun启动torchrun --nproc_per_node4 train_multi_GPU.py --lr 0.004 --amp--nproc_per_node为使用 GPU 数量若只想使用部分 GPU 设备可在指令前加上CUDA_VISIBLE_DEVICES环境变量指定设备编号例如只使用物理设备第 1 块与第 4 块CUDA_VISIBLE_DEVICES0,3 torchrun --nproc_per_node2 train_multi_GPU.pytrain_multi_GPU.py 与单机版的主要差异通过init_distributed_modetrain_utils/distributed_utils.py读取RANK / WORLD_SIZE / LOCAL_RANK环境变量完成分布式初始化使用DistributedSampler划分数据DistributedDataParallel包装模型可选--sync-bn开启同步 BN多卡间同步训练会变慢参数略有扩展-j / --workers默认 4、--output-dir默认./multi_train、--world-size、--dist-url默认env://、--test-only仅测试不训练等其中--batch-size表示每块 GPU 上的 batch size全局 batch size 为NGPU × batch_size日志写入与权重保存仅在 rank 0 主进程执行save_on_master。注意README 注明 Windows 暂不支持多 GPU 训练torchrun与nccl后端依赖 Linux 环境多 GPU 训练建议在 Ubuntu / CentOS 上进行。八、训练记录与结果对比README 记录了u2net_full在本仓库配置下的训练结果。训练指令4 卡、学习率 0.004、开启 AMPtorchrun --nproc_per_node4 train_multi_GPU.py --lr 0.004 --amp最终在 DUTS-TE 上的验证结果MAE: 0.047 maxF1: 0.859与官方权重验证结果MAE 0.044 / maxF1 0.868相比略有差距符合正常复现水平差异主要来自 post_norm 后处理与训练细节。训练过程详情可参考仓库中的 results20220723-123632.txt从该记录可以看到典型的收敛曲线第 0 个 epochtrain_loss: 2.7385 lr: 0.002002 MAE: 0.465 maxF1: 0.464warmup 初期第 100 个 epochMAE: 0.055 maxF1: 0.847第 190 个 epoch 附近进入平台期MAE: 0.047 maxF1: 0.860最终第 359 个 epochtrain_loss: 0.1518 lr: 0.000000 MAE: 0.047 maxF1: 0.859。整个训练过程学习率从 warmup 后约 0.004 沿余弦曲线平滑衰减到 0MAE 从 0.465 稳步降至 0.047 左右maxF1 从 0.464 提升至约 0.859验证集指标在后期趋于稳定——这是判断训练收敛与是否需要早停的直接依据。九、验证与推理9.1 单独验证validation.py使用 validation.py 对指定权重在 DUTS-TE 上单独验证python validation.py --data-path /path/to/DUTS_root --weights ./u2net_full.pth参数说明--data-path为 DUTS 根目录--weights指定权重文件路径--device默认cuda:0。脚本构建SODPresetEval([320, 320])预处理后以 batch_size1 逐张送入网络由evaluate计算 MAE 与 maxF1 并打印。权重加载兼容两种格式直接 load_state_dict或从训练保存的含model键的 checkpoint 中提取validation.py。9.2 单图预测predict.pypredict.py 提供了最简推理流程weights_path ./u2net_full.pth img_path ./test.png threshold 0.5执行python predict.py前将weights_path指向权重、img_path指向测试图片即可。其推理流程要点图片 BGR→RGB 后依次做ToTensor → Resize(320) → NormalizeImageNet 均值/方差加载u2net_full()权重并model.eval()先用全零init_img预热模型触发 CUDA kernel 初始化再对真实图片推理并统计耗时推理输出torch.sigmoid后的单通道显著图缩放到原图尺寸按threshold0.5二值化生成 Mask用显著区域origin_img * pred_mask提取前景并保存pred_result.png同时用 matplotlib 显示。9.3 ONNX 导出src/model.py 末尾还附带convert_onnx函数可将 RSU / RSU4F 或完整 U2-Net 以opset_version11导出为 ONNX输入尺寸 288×288方便在部署侧使用该仓库另有独立的模型部署章节可供参考。十、总结与常见注意事项综合 README 与源码使用本仓库训练/使用 U2-Net 时需要重点注意以下几点任务定位U2-Net 属于显著性目标检测SOD标签为二值 Mask评估指标是 MAE 与 maxF1与语义分割任务不可混用数据集结构--data-path必须指向包含DUTS-TR与DUTS-TE的根目录且训练图片为.jpg、Mask 为同名.png官方权重加载官方权重包含conv.bias键而本仓库ConvBNReLU默认biasFalse可用 convert_weight.py 转换权重或按 README 说明将 bias 改为 True指标差异来源本仓库验证时移除了 post_norm(output - mi) / (ma - mi)导致 maxF1 相比原论文约低 0.0040.868 vs 0.872如需复现原论文数值可自行恢复该后处理多 GPU 训练使用torchrun --nproc_per_nodeN train_multi_GPU.py配合CUDA_VISIBLE_DEVICES指定 GPU仅支持 LinuxWindows 不支持多 GPU验证 batch_sizeMAE / F1 指标实现要求验证时 batch_size 必须为 1脚本已固定无需修改深度监督损失训练时融合输出与 6 个侧输出各计算一次 BCE 并求和这是 U2-Net 多尺度学习能力的关键也是 src/model.py 中self.training分支存在的意义。至此从网络结构、数据集、训练、验证到推理的完整链路已全部打通。若需进一步深入理解 U2-Net 代码细节与训练策略可直接阅读 src/model.py、train_utils/train_and_eval.py 与 train_utils/distributed_utils.py 的完整实现并参照仓库内results*.txt复现训练过程。赞分享示例工程【免费下载链接】deep-learning-for-image-processingdeep learning for image processing including classification and object-detection etc.项目地址https://gitcode.com/gh_mirrors/de/deep-learning-for-image-processing点击查看免费下载相关推荐U-Net 图像分割实战基于 deep-learning-for-image-processing 仓库的 DRIVE 视网膜血管分割与 PyTorch 训练部署指南U Net 图像分割实战基于 deep learning for image processing 仓库的 DRIVE 视网膜血管分割与 PyTorch 训练示例工程Colorify.js扩展开发指南如何自定义插件和增强颜色提取功能Colorify.js扩展开发指南如何自定义插件和增强颜色提取功能 Colorify.js是一个简单、可定制、轻量级的JavaScript颜色提取库能够从图Mask R-CNN 实例分割实战基于 deep-learning-for-image-processing 仓库的 COCO / Pascal VOC 训练、验证与部署指南Mask R CNN 实例分割实战基于 deep learning for image processing 仓库的 COCO / Pascal VOC 训练示例工程上一篇CloddsBot多语言设置10种语言i18n切换完整教程下一篇Sketch MeaXure重新定义设计标注工作流的现代解决方案创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考