跳至正文
仇天赐.研究与工程笔记
搜索标题、摘要和标签 Esc 关闭
返回技术档案

把 U-Net 基线做扎实:先让数据与训练闭环

围绕输入契约、标签检查、二分类损失与完整推理流程,整理一个可诊断、可复现的 U-Net 分割基线。

回顾整理

本文归入 2025.02 主题档案,整理于 2026.11.05。月份用于组织阅读路径,不代表具体实验或项目的完成日期。

基线要回答哪些问题

这篇方法回顾按主题归档,不表示在归档月份开展了特定实验。建立分割基线的目的,是确认数据、标签、优化和推理能否形成可信闭环,让后续改动有可比较的参照。

U-Net 通过收缩路径获取上下文,并通过扩张路径恢复定位信息,是合适的起点。U-Net 原论文 这里讨论的是围绕该结构的工程方法,不声称复现论文的全部实现,更不提供未经测量的成绩。

先把一个样本完整看明白

先列出输入维度、通道含义、单位、空间方向、类别编号和忽略标签。图像与掩膜不只要形状相同,还要落在同一空间位置。重采样、裁剪、翻转后,都应叠加检查边缘。

二分类任务可以约定图像为 [B, C, H, W],标签为 [B, 1, H, W],标签只含 0 和 1。标签缩放不能使用会制造小数类别的连续插值;几何增强应同步作用于图像和标签,强度增强则只作用于图像。

选择 2D、相邻切片输入或 3D 时,应考虑目标连续性、体素各向异性和资源预算。无论哪种输入,同一受试者的数据都应遵守统一的集合划分。

损失函数先服务于诊断

对单通道二分类输出,可组合二元交叉熵与软 Dice 损失。下面只是一个明确约定的起点,要求标签无忽略区域且已转换为 0/1。

import torch
import torch.nn.functional as F

def binary_loss(logits, target, eps=1e-6):
    if logits.ndim < 3 or logits.shape[1] != 1:
        raise ValueError("expected a single-channel segmentation batch")
    if logits.shape != target.shape:
        raise ValueError("prediction and target must align")
    target = target.to(device=logits.device, dtype=logits.dtype)
    bce = F.binary_cross_entropy_with_logits(logits, target)
    prob = torch.sigmoid(logits)
    dims = tuple(range(1, logits.ndim))
    intersection = (prob * target).sum(dim=dims)
    total = prob.sum(dim=dims) + target.sum(dim=dims)
    dice = (2 * intersection + eps) / (total + eps)
    return bce + (1 - dice).mean()

交叉熵分支直接接收 logits,不要先做 sigmoid 再传入对应接口;组合接口使用更稳定的计算形式。PyTorch BCEWithLogitsLoss

软 Dice 的平滑项属于训练定义,不能直接替代评估时的空掩膜政策。多类别互斥任务还需重新定义输出、标签和损失,不能机械扩展这段代码。

用小规模过拟合检查管线

在进行正式训练前,可以暂时关闭随机增强,让模型拟合一个很小的固定训练子集。这里只检查管线是否具备学习能力,不用该子集估计泛化表现。

若损失下降而叠加图不对,应先查标签方向、输出通道和可视化逻辑;若始终输出背景,应查前景采样、类别映射与梯度。通过这一步以后,再恢复增强、正常采样和独立验证。不要用增加网络复杂度掩盖数据错误。

前景采样可以帮助模型看到稀少目标,但应保留空背景样本,让训练覆盖误报场景。采样器的实际输出值得单独审阅:声明要采前景,并不保证裁剪之后仍包含有效目标。增强后的标签若变空,也要区分正常现象与坐标错误。

若数据包含未标注区域,不能把它们自动当作背景。此时需要有效区域掩膜,并在损失的分子、分母和交叉熵归约中一致排除;上述简化函数并不覆盖这种标签定义。

推理闭环与失败分析

验证应覆盖完整图像或体数据,并固定重叠窗口融合、阈值和坐标恢复方式。只在含目标的裁剪块上评估,会漏掉完整视野中的误报。最大连通域等后处理也可能删除真实多发目标,必须作为独立配置记录。

对照实验每次只调整一个主要因素,查看患者级指标与失败样本。若前景 Dice 看似稳定,却总漏掉细长结构,下一步应研究分辨率、采样与边界,而不是只追求更低训练损失。nnU-Net 的启发之一,是把数据适配、训练和推理当成共同设计的系统。nnU-Net,2018

参考

笔记至此,探索继续。交流与纠错