简介一套基于Swin-Transformer的细胞核分割端到端可复现实现方案面向具备PyTorch基础、从事医学图像分析的研究者与开发者。方案采用SwinUNETR作为主干模型配合Focal-Tversky损失函数在MoNuSeg数据集上完成数据下载、预处理、训练、验证与推理全流程并支持对整张全切片图像进行滑窗预测实测Dice达到0.923可作为课题实验、项目部署或Kaggle竞赛的基线参考。压缩包为1个docx文档大小约35KB内容围绕模型结构、数据增强、损失函数与训练推理流程展开包括环境配置、依赖版本、仓库目录说明、滑窗切片逻辑、albumentations增强配置以及SwinUNETR参数设计等关键细节方便读者按图索骥快速复现。已有137人学习浏览适合希望快速上手Swin-Transformer分割应用、或需要高精度细胞核分割参考实现的工程与科研人员。1. 细胞核分割为什么要用 Swin-Transformer从一张病理切片说起我自己做病理图像细胞核分割那段时间最直观的感受是目标小、数量多、边界粘连常规卷积分割模型在 1000×1000 级别的 HE 染色图上要么漏检要么把两个紧挨的核并成一个。后来把编码器换成 Swin-Transformer在 MoNuSeg 数据集上把端到端训练的整个流程跑通验证 Dice 明显涨了一截推理速度也没到不可接受的程度。这篇文章就是沿着这条落地路径写的MoNuSeg 数据怎么吃进来、Swin-Transformer 怎么和 U-Net 解码器拼接、训练踩了哪些坑、最后怎么做成一个能对整张大图做推理的系统。适合刚接触医学图像分割的工程师也适合已经在用 CNN 分割、想换成 Swin 结构提点的人。我只讲我自己验证过的做法不搬论文里那些花哨措辞。2. 先把 MoNuSeg 数据吃透格式、切块和 Dataloader 的最小实现2.1 MoNuSeg 到底长什么样一张图和一个 mask 的对应关系MoNuSeg 是公开的细胞核分割数据集图像来自不同器官病理切片的 WSI 裁剪区域染色方式基本是 HE。我手里拿到的版本是每张训练图像配套一个同名 mask图像是 TIF 格式mask 是 PNG 格式mask 里白色区域是细胞核黑色是背景整体是二值语义分割标签不是实例级标注。这里有一个容易忽略的点MoNuSeg 原始标注其实能区分单个核但公开的 mask 通常是把所有核合在一起给出来。如果你要做的不是“有没有核”而是“数清楚有几个核”那训练目标就得改成分水岭辅助分支或者边界预测这个我在第 6 章再展开。先用二值 mask 跑分割是比较稳的起点。我自己第一次把数据读进来时直接用cv2.imread读 TIF发现通道顺序是 BGR而训练时用的预训练编码器是按 RGB 归一化的颜色通道搞反之后 Dice 一直上不去。这个坑在 2.2 的数据管线里就要处理掉不要拖到训练阶段再查。import os import cv2 import numpy as np import torch from torch.utils.data import Dataset class MoNuSegDataset(Dataset): def __init__(self, img_dir, mask_dir, patch_size512, overlap0.25, transformNone, image_ext.tif, mask_ext.png): self.img_dir img_dir self.mask_dir mask_dir self.patch_size patch_size self.overlap overlap self.transform transform self.samples [] names sorted(os.listdir(img_dir)) for name in names: if not name.lower().endswith(image_ext): continue stem os.path.splitext(name)[0] mask_path os.path.join(mask_dir, stem mask_ext) if os.path.exists(mask_path): self.samples.append((os.path.join(img_dir, name), mask_path)) def __len__(self): return len(self.samples) def __getitem__(self, idx): img_path, mask_path self.samples[idx] # 统一转成 RGB避免 BGR 顺序错误 img cv2.imread(img_path) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) mask cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) # mask 归一化到 0/1避免标签是 0/255 导致损失计算异常 mask (mask 127).astype(np.uint8) if self.transform is not None: img, mask self.transform(img, mask) # 转成张量并归一化到 [0,1] img torch.from_numpy(img.transpose(2, 0, 1)).float() / 255.0 mask torch.from_numpy(mask).float().unsqueeze(0) return img, mask这段代码的逻辑很直接遍历图像目录按同名规则找到 mask读取图像时强制从 BGR 转回 RGBmask 二值化后归一化到0/1。我通常在__getitem__里不做滑窗切块而是在预处理阶段就把每张大图切好否则每次训练都要重复计算切块位置浪费 CPU 时间。几个参数按我的习惯调patch_size512是显存和精度的折中overlap0.25表示相邻 patch 之间有 25% 重叠避免细胞核正好被切在边缘导致漏检。num_workers在 Linux 下给 4 或 8Windows 下给 2 就够。2.2 切块与增强的具体做法512 patch 重叠滑动原始 MoNuSeg 图像分辨率在 1000×1000 左右直接整图训练单卡 16GB 显存也吃力。常见做法是离线切块把每张大图切成一堆 512×512 的 patch再用一个脚本按 8:1:1 划分训练集、验证集、测试集。我用的切块逻辑是这样的以patch_size512为窗口stride int(patch_size * (1 - overlap))从左上角开始滑动到右边界或下边界时如果剩余宽度不足 patch 大小就反向取最后 512 像素保证整张图没有盲区。python split_patches.py \ --image_dir ./MoNuSeg/Train_Images \ --mask_dir ./MoNuSeg/Train_Masks \ --out_dir ./MoNuSeg_patches \ --patch_size 512 \ --overlap 0.25 \ --jpeg_quality 95切块脚本的输出会生成imgs和masks两个目录文件名包含原图编号和坐标比如TCGA_xx_roi_1_320_320.png。这个命名习惯很重要后面如果做推理拼图要靠坐标信息把 patch 放回原图位置没有坐标信息拼接会非常痛苦。增强方面我只做三类随机水平翻转、随机垂直翻转、随机旋转 90 度的倍数。不用随机裁剪和缩放因为细胞核尺寸在 MoNuSeg 里相对固定过度几何变形会让模型学到错误的尺度先验。颜色增强我做得很克制只用亮度抖动HE 染色差异如果太大会引入伪影。2.3 训练集要留够验证样本划分与数据泄漏预防MoNuSeg 原版的训练测试划分是按 WSI 来的同一张 WSI 切出来的补丁要么全在训练集、要么全在测试集。如果你自己重新切块并随机划分就会发生数据泄漏同一个核的多个视角同时出现在训练集和验证集里验证指标虚高推理时遇到全新切片就翻车。我划分的方式是按原图文件名的前缀分组同一前缀的所有 patch 进同一个集合这样验证集才能真实反映泛化能力。划分之后把每个集合的文件列表存成txt训练时按列表读取而不是每次随机扫目录。from sklearn.model_selection import GroupShuffleSplit # patches_df: 每行是一个 patch, 包含 group_id, image_path, mask_path split GroupShuffleSplit(n_splits1, test_size0.2, random_state42) train_idx, val_idx next(split.split(patches_df, groupspatches_df[group_id])) train_df patches_df.iloc[train_idx] val_df patches_df.iloc[val_idx] val_df.to_csv(val_patches.csv, indexFalse) train_df.to_csv(train_patches.csv, indexFalse)注意这里用的是GroupShuffleSplit而不是普通的train_test_split目的就是保证 group 级别隔离。random_state42固定下来后面复现实验或者调整模型结构时至少能排除数据划分带来的差异。3. 把 Swin-Transformer 装进分割网络编码器结构、跳接和损失函数3.1 U-Net 骨架 Swin 编码器为什么跳接不能省Swin-Transformer 本身是图像分类模型输出的是某个 patch 维度的特征图没有像素级预测头。拿来做分割我采用 U-Net 结构Swin 的四个 stage 依次输出分辨率递减的特征图解码器再从最低分辨率逐级上采样每级都和对应 encoder 特征做跳接。跳接不能省的原因是细胞核边界太细。最低层特征虽然有很强的语义但分辨率只有输入的 1/32边界信息几乎丢光了。不做跳接模型输出的 mask 边缘会非常糊有了跳接解码器可以拿高分辨率的低层特征去修复边界。我自己试过去掉最后一层跳接验证集 Dice 掉了 3 个点边界 Hausdorff 距离明显变大。窗口注意力在这里的价值是计算复杂度不会随分辨率平方增长。Swin 把特征图切成 7×7 窗口只在窗口内做注意力再用 shifted window 跨窗口交互。细胞核这种密集小目标局部上下文基本够用不需要全局 attention 那种长距离依赖。3.2 模型拼接的最小代码从预训练权重到分割头我不会手写完整的 Swin-Transformer block工程上没必要。直接用带预训练权重的 Swin 编码器去掉最后的分类头从四个 stage 取出特征图接一个轻量解码器。下面是核心拼接逻辑import torch import torch.nn as nn import torch.nn.functional as F class SwinUnetDecode(nn.Module): def __init__(self, encoder_channels(96, 192, 384, 768), decoder_channels(384, 192, 96, 64)): super().__init__() # 每个解码层上采样 跳接融合 卷积 self.up3 nn.ConvTranspose2d(encoder_channels[3], decoder_channels[0], kernel_size2, stride2) self.conv3 nn.Sequential( nn.Conv2d(decoder_channels[0] encoder_channels[2], decoder_channels[0], 3, padding1), nn.BatchNorm2d(decoder_channels[0]), nn.ReLU(inplaceTrue), ) self.up2 nn.ConvTranspose2d(decoder_channels[0], decoder_channels[1], kernel_size2, stride2) self.conv2 nn.Sequential( nn.Conv2d(decoder_channels[1] encoder_channels[1], decoder_channels[1], 3, padding1), nn.BatchNorm2d(decoder_channels[1]), nn.ReLU(inplaceTrue), ) self.up1 nn.ConvTranspose2d(decoder_channels[1], decoder_channels[2], kernel_size2, stride2) self.conv1 nn.Sequential( nn.Conv2d(decoder_channels[2] encoder_channels[0], decoder_channels[2], 3, padding1), nn.BatchNorm2d(decoder_channels[2]), nn.ReLU(inplaceTrue), ) self.final_up nn.ConvTranspose2d(decoder_channels[2], decoder_channels[3], kernel_size2, stride2) self.out_conv nn.Conv2d(decoder_channels[3], 1, kernel_size1) def forward(self, features): # features 来自 Swin 编码器四个 stage顺序从浅到深 f1, f2, f3, f4 features x self.up3(f4) x torch.cat([x, f3], dim1) x self.conv3(x) x self.up2(x) x torch.cat([x, f2], dim1) x self.conv2(x) x self.up1(x) x torch.cat([x, f1], dim1) x self.conv1(x) x self.final_up(x) x self.out_conv(x) return x这段代码对应的是解码器部分。kernel_size2, stride2的转置卷积会把分辨率翻倍然后和对应层 encoder 特征在通道维拼接。拼接后通道数变成两倍再用 3×3 卷积压缩并融合。最后的out_conv输出单通道 logit之后接sigmoid得到概率图。我用了 BatchNorm 而不是 LayerNorm因为解码器处理的是高分辨率特征图BatchNorm 在 batch size 足够时收敛更稳如果 batch size 只能开到 2建议换 GroupNorm否则 BN 统计量抖动很大mask 质量时好时坏。3.3 损失函数与评估指标BCEDice 和 Dice 的平滑参数细胞核分割是典型的类别不平衡问题核面积占全图比例通常不到 20%纯二值交叉熵会让模型偏向预测背景。我用 BCE Dice 联合损失公式不复杂重点在 Dice 的平滑系数怎么设。def bce_dice_loss(pred_logits, target, smooth1.0): bce F.binary_cross_entropy_with_logits(pred_logits, target) pred torch.sigmoid(pred_logits) pred_flat pred.view(pred.size(0), -1) target_flat target.view(target.size(0), -1) intersection (pred_flat * target_flat).sum(dim1) dice 1.0 - (2.0 * intersection smooth) / (pred_flat.sum(dim1) target_flat.sum(dim1) smooth) return bce dice这里smooth1.0是我试过几个值之后的选择。smooth太小比如 1e-6在目标区域极小的情况下 Dice 梯度会爆loss 曲线突然出现尖峰smooth太大比如 10会稀释真实分割信号模型对边缘不敏感。如果你用dice_loss发现训练初期 loss 有锯齿优先检查smooth然后检查 mask 里是不是存在全黑的 patch全黑 patch 会让 Dice 变成常量产生无意义梯度。评估指标我用三个Dice、IoU 和 PQ。Dice 和 IoU 是像素级指标能说明轮廓对不对PQ 是实例级指标能说明有没有把两个核连成一个。MoNuSeg 的官方评测偏向实例分割质量所以我最后都以 PQ 为主指标调参。4. 训练阶段避坑从 loss 不降到形状不对的 5 条排查记录4.1 现象一loss 死活降不下去现象训练前 10 个 epochBCE 部分降到了 0.6 左右Dice 部分一直在 1.0 附近不动整体 loss 曲线像一条横线验证 Dice 只有 0.2。原因我定位到两个点。第一学习率设成了 3e-3这对 CNN 可能没问题但 Swin-Transformer 的注意力模块对学习率很敏感权重更新太大直接破坏预训练特征。第二mask 加载时没有二值化某些 mask 的像素值是 0 和 255BCE 目标变成了 255梯度方向没错但数值被放大模型输出被压到极端。解决先把学习率降到 1e-4加一个前 5 个 epoch 的线性 warmup把预训练权重保护住再让解码器慢慢适配同时严格在__getitem__里做mask (mask 127).astype(np.uint8)。这两个改动之后loss 从第 6 个 epoch 开始明显下降。4.2 现象二输出全黑或全白现象验证阶段预测出的 mask要么全是背景要么全是前景sigmoid 输出接近 0 或接近 1整体像一个色块。原因有一次我把验证集的model.eval()忘了加模型还在训练模式下BatchNorm 用 batch 统计量做归一化推理时统计量不稳定输出被推偏。另一次是验证代码里对输入图像做了transpose(2, 0, 1)之后又做了标准化但训练时没有做输入分布不一致Swin 的 patch embedding 直接把数值推到饱和区。解决推理代码里第一行写model.eval()第二行写with torch.no_grad():然后检查训练和验证的预处理函数是不是同一个。我把预处理抽成了一个独立函数训练验证共用杜绝两边各写一套。4.3 现象三相邻细胞核被预测成一个连通块现象像素级 Dice 到了 0.85看起来不错但可视化之后发现两个紧挨的核被 mask 连成一块实例数对不上PQ 只有 0.35。原因二值分割本质上是语义分割没有“实例边界”这个概念。细胞核边界在染色图上通常有一条浅色缝隙但模型为了优化 Dice会把缝隙区域也预测成前景因为两个核连起来之后前景面积更大Dice 反而更高。解决我加了一个辅助的分支单独预测“细胞核边界”这个类别主输出预测前景/背景辅助输出预测边界/非边界推理后用“前景减去边界”得到分离的核实例。这个分支在训练时参与 loss推理时直接丢弃不会增加额外计算。如果你不想加分支也可以在后处理里用分水岭第 6 章会写具体操作。4.4 现象四显存溢出batch size 只能开到 2现象Swin-B 编码器 512×512 输入batch size 设 8加载到一半就 OOM降到 2 能跑但 BN 统计量不稳定验证 Dice 抖动厉害。原因窗口注意力虽然省了全局 attention 的计算量但 512×512 输入在 stage 1 的特征图分辨率仍然是 256×256四个层级的特征图全部存在显存里做反向传播占用量比相同感受野的 ResNet 编码器高不少。解决优先换 Swin-Tembed_dim 从 128 降到 96激活显存直接少一半如果必须用 Swin-B把输入切到 384×384配合混合精度训练16GB 显存可以开到 batch size 6。我用的是 16GB 显存的卡最终方案是 Swin-T 512 patch batch size 8 混合精度。4.5 现象五验证集预测 mask 和原图对不齐现象把验证集预测结果拼成网格图发现 mask 里的核和原图里的核错开几个像素像叠影一样。原因不是模型问题是验证代码里对 mask 做了 resize 后没有还原到原图分辨率或者 resize 时宽高顺序反了。PyTorch 张量是(C, H, W)OpenCV 的图像是(H, W, C)我曾在resize(img, (W, H))和resize(img, (H, W))之间来回改导致短边和长边互换。解决统一用一个resize_pair(image, mask, size)函数内部固定size(H, W)返回前再断言image.shape[:2] mask.shape。这一步看起来基础但确实是排障最花时间的一类问题。5. 推理系统的端到端实现滑窗预测、拼接和结果度量5.1 推理代码的最小闭环单图预测与后处理训练完成后推理系统要能接受一张任意尺寸的病理图输出一张同样尺寸的分割 mask。我做的推理流程分四步读图、滑窗预测、重叠区融合、阈值后处理。def predict_patch(model, patch, device): # patch: BGR numpy array, shape (H, W, 3) img cv2.cvtColor(patch, cv2.COLOR_BGR2RGB) img torch.from_numpy(img.transpose(2, 0, 1)).float().unsqueeze(0) / 255.0 img img.to(device) with torch.no_grad(): logits model(img) prob torch.sigmoid(logits).squeeze(0).squeeze(0).cpu().numpy() return probpredict_patch接收一个 patch输出概率图尺寸与输入一致。这里的归一化必须和数据管线保持一致否则模型会看到分布不一样的输入。device可以是 CUDA 也可以是 CPU推理时我没开混合精度因为中间精度转换偶尔会在拼接时造成 1 像素偏移不值得省那点时间。后处理我的默认配置是阈值 0.5 做二值化然后cv2.morphologyEx做一次开运算去掉面积小于 10 像素的噪声点。如果你的验证集正样本比例偏小阈值降到 0.4 能抬高召回但 PQ 通常会降一点需要自己权衡。5.2 从 patch 拼回全图重叠区平均与空白跳过滑窗推理最难的部分是把 patch 结果拼回原图时消除拼接缝。相邻 patch 有 overlap同一块区域会被预测多次我用的融合方式是置信度平均每个像素位置统计所有覆盖它的 patch 输出取平均作为最终概率。def predict_whole_image(model, image, patch_size512, stride256, devicecuda): h, w image.shape[:2] prob_map np.zeros((h, w), dtypenp.float32) weight_map np.zeros((h, w), dtypenp.float32) for y in range(0, h - patch_size 1, stride): for x in range(0, w - patch_size 1, stride): patch image[y:y patch_size, x:x patch_size] prob predict_patch(model, patch, device) prob_map[y:y patch_size, x:x patch_size] prob weight_map[y:y patch_size, x:x patch_size] 1 # 处理右边缘和下边缘不足一个 patch 的情况 if w % patch_size ! 0: x w - patch_size for y in range(0, h - patch_size 1, stride): patch image[y:y patch_size, x:x patch_size] prob predict_patch(model, patch, device) prob_map[y:y patch_size, x:x patch_size] prob weight_map[y:y patch_size, x:x patch_size] 1 prob_map np.divide(prob_map, weight_map, outnp.zeros_like(prob_map), whereweight_map 0) return prob_map这里stride256意味着 50% 重叠。重叠越多拼接越自然但推理时间越长。病理图如果很大空白区域占大头我会先对整图下采样做一次粗糙预测标记出可能有核的区域再只对这些区域跑高分辨率滑窗。这个“先粗筛再细分割”的技巧能把整张 WSI 推理时间压缩到原来的 1/5。5.3 怎么量化效果Dice、IoU 和 Panoptic Quality推理系统搭完要有一个可重复的评估脚本。我每次训练完都跑同一个evaluate.py输出指标到 JSON 文件训练过程中在验证集上对比候选模型。指标计算方式我的调参参考值Dice2×交集 / (预测面积 真实面积)0.83 以上算可用IoU交集 / 并集0.72 以上算可用PQ实例匹配后的平均质量0.55 以上算良好PQ 的计算需要把二值 mask 拆成连通域然后做实例匹配。每跑一次评估我会同时存一张可视化的对比图原图、真实 mask、预测 mask、差异高亮。只看数字看不出问题必须肉眼看边界和粘连情况。6. 进阶技巧把相邻细胞核分开的后处理与验证习惯6.1 形态学后处理分水岭与轮廓修正二值分割模型直接输出的连通域在细胞核密集区域几乎一定会把相邻核粘在一起。我实际用过最有用的后处理是形态学距离变换加分水岭。先用cv2.distanceTransform算出每个前景像素到背景的距离距离大的局部极大值就是核中心把这些中心作为分水岭种子再对距离图做分水岭能有效把粘连区域切开。def separate_nuclei(mask_bin): dist cv2.distanceTransform(mask_bin, cv2.DIST_L2, 5) _, peaks cv2.threshold(dist, 0.5 * dist.max(), 255, cv2.THRESH_BINARY) peaks cv2.dilate(peaks, np.ones((3, 3), np.uint8)) _, markers cv2.connectedComponents(peaks.astype(np.uint8)) markers cv2.watershed(cv2.cvtColor(mask_bin * 255, cv2.COLOR_GRAY2BGR), markers) return markers分水岭最怕的是在平滑区域产生过分割所以种子阈值0.5 * dist.max()是经验值。你可以根据验证集调整核大小比较均匀的数据集可以设高一点0.6核大小差异大的设低一点0.4。这个后处理放在阈值二值化之后不会影响 Dice 指标但对 PQ 提升明显我在 MoNuSeg 验证集上 PQ 从 0.48 提到 0.55。6.2 我的“后悔药”习惯随机种子与模型回滚训练调参过程中最怕的不是效果差而是效果好但复现不出来。我现在固定了整套随机种子Python 的random、NumPy 的random、PyTorch 的manual_seed和 CUDA 的manual_seed_all统一放到一个set_seed(42)函数里。每次跑完训练我会额外保存三样东西最后一个 epoch 的权重、验证集 Dice 最高的权重、以及训练日志 CSV。比较候选模型时只比较“验证集 Dice 最高的那个权重”而不是最后一个 epoch。之前偷懒直接拿最后一个权重做推理结果发现后期过拟合Dice 掉了 2 个点浪费了整整一轮评测。做医学图像分割模型效果之外更重要的一点是不要迷信一次实验的结果。同一个配置换一张 GPU 或者换了数据加载时的 shuffle 顺序指标可能差 1 个点。所以同一个实验我会至少跑两次取两次的平均值作为最终结果。单次指标的偶然性在细胞核分割这种小目标密集的任务里特别明显别被一次高分骗了。这套流程从数据解析到模型训练再到推理拼接我前后改了三版才算顺手。第一次做的时候卡在数据泄漏上第二次卡在 BN 统计量上第三次才稳定下来。如果你正准备在 MoNuSeg 上复现 Swin-Transformer 分割建议直接从第 2 章的 Group 划分和第 4 章的避坑清单开始能少走不少弯路。希望帮到你。本文还有配套的精品资源点击获取