简介面向医学图像分割研究者与深度学习开发者这份基于TransUnet的DRIVE视网膜血管分割实战资源包含了完整可运行代码和已标注背景/前景的DRIVE数据集可端到端完成从模型训练到评估推理的全流程。压缩包共76个文件整体仅7.87MB其中40个png文件用于数据集样本与分割结果可视化18个py源码文件覆盖训练、验证、推理三大核心脚本并包含Transformer、Unet、vit等关键模块15个pyc为运行缓存另附requirements.txt、README与txt说明文档借助目录结构和注释可快速定位各功能模块。目前已有323人浏览学习适合医学图像领域入门及进阶者。代码注释非常详细train脚本会生成训练集和验证集的loss、IoU曲线、学习率衰减曲线、训练日志以及数据集可视化图像evaluate脚本可计算测试集的IoU、recall、precision、像素准确率等评估指标predict脚本支持单张图像推理输出GT掩膜与GT叠加原图的融合图便于直观检验分割效果。配套README给出了傻瓜式训练步骤按说明修改路径即可迁移到自定义数据实用性强。1. TransUnet 在 DRIVE 上的分割实战拿到开源项目后到底该怎么跑通、怎么改、怎么评估眼底血管分割是医学影像分割里最经典的入门场景DRIVE 数据集也常被拿来当作算法验证的基准。这个项目给的是一套基于 TransUnet 的完整实现不是只有模型结构而是把训练、评估、推理三个环节全做了闭环。你拿到手里会发现train.py 负责产出训练日志和曲线evaluate.py 负责算测试集指标predict.py 负责出掩膜图三件事分得清清楚楚。更友好的一点是代码里做了详细注释README 里写了训练自定义数据的流程基本能称得上傻瓜式运行。适合三类人刚接触语义分割、想用 Transformer 结构跑通一条完整流程的研究生需要拿 DRIVE 做 baseline 对比的算法工程师以及想快速迁移到自有数据集、但不想从零搭框架的从业者。2. 项目解包文件清单与三个脚本的分工逻辑2.1 压缩包里到底有什么解开压缩包后你会看到两套核心目录data/下存放数据集相关结构和代码unet/下则是模型定义。先把这个结构理清楚比急着跑 train.py 重要得多。Transformer-Unet/ ├── data/ │ ├── train/ # 训练集数据按 DRIVE 原始划分存放 │ ├── val/ # 验证集数据 │ └── test/ # 测试集数据 ├── unet/ │ ├── __init__.py │ ├── vanilla_transformer.py # 原始 Transformer 模块 │ ├── unet_transformer.py # TransUnet 的核心编码器-解码器结构 │ └── common.py # 公共组件如 patch embedding ├── train.py # 训练入口 ├── evaluate.py # 测试集评估入口 ├── predict.py # 推理入口 ├── utils.py # 工具函数 ├── dataset.py # 数据加载与预处理 ├── confuse_matrix.py # 混淆矩阵与指标计算 ├── requirements.txt # 依赖清单 └── README.md # 使用说明文件层级不深但每个文件的责任边界很清楚。第一次接触这个项目时建议先读 README 和 dataset.py再去看模型定义顺序反了容易一头扎进 Transformer 的细节里出不来。2.2 三个入口脚本分别承担什么职责这个项目最大的特点是把训练、评估、推理拆成了三个独立脚本不像很多开源项目把全部逻辑塞进一个 main 函数。这种拆法在实际工作中更实用——训练和评估解耦后你可以随时用同一套权重去验证不同的测试集不需要重新训练。train.py 做的事情不只是迭代数据更新参数它会在训练过程中持续产出可视化和日志文件训练集和验证集的 loss 曲线、IoU 曲线、学习率衰减曲线以及数据集的原始图像和标注可视化图。这意味着训练结束后你不需要额外写代码就能复盘整个训练过程。confuse_matrix.py 是评估阶段的核心依赖。evaluate.py 调用它计算测试集上的 IoU、Recall、Precision、像素准确率等指标最终输出的不是一两个数字而是一套完整的混淆矩阵指标。这些指标对医学影像分割尤其重要因为背景像素占绝对多数只看准确率会严重虚高。predict.py 负责最后的推理。它会读入测试图像输出两个结果gt_mask真实标注掩膜和预测的掩膜同时生成将预测叠加在原图上的可视化图像。这个设计很贴心因为分割任务的效果好不好肉眼看叠加图往往比看数值更直观。2.3 依赖环境与配置要点requirements.txt 里列出了项目运行所需的库。从实际经验看最核心的依赖集中在 PyTorch、TorchVision、NumPy、OpenCV 和 Matplotlib 上。如果你之前跑过其他深度学习项目大概率不需要重新安装太多东西。需要特别注意的是 Python 版本。这个项目里出现了dataset.cpython-310.pyc文件说明作者是在 Python 3.10 环境下打包的。你本地最好也用 3.10 或更高版本避免因为 Python 版本差异导致部分语法或库接口不兼容。CUDA 环境建议用 11.x 以上版本显存至少 8GB因为 TransUnet 的参数量比纯 CNN 结构大不少。3. 跑通训练全流程从数据加载到指标曲线3.1 DRIVE 数据集的加载与预处理逻辑dataset.py 负责把磁盘上的图像读入内存并做预处理。这里最关键的是搞清楚 DRIVE 数据集的格式它包含训练集、验证集和测试集每张图都有对应的血管标注图。DRIVE 的原始图像是 565×584 的 RGB 眼底图标注图是单通道二值图。# dataset.py 中核心的数据加载逻辑简化版 import cv2 import numpy as np import torch from torch.utils.data import Dataset class DriveDataset(Dataset): def __init__(self, img_dir, mask_dir, img_size512, transformNone): self.img_paths sorted(glob.glob(os.path.join(img_dir, *.png))) self.mask_paths sorted(glob.glob(os.path.join(mask_dir, *.png))) self.img_size img_size self.transform transform def __len__(self): return len(self.img_paths) def __getitem__(self, idx): img cv2.imread(self.img_paths[idx]) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) mask cv2.imread(self.mask_paths[idx], cv2.IMREAD_GRAYSCALE) # 统一缩放到模型输入尺寸 img cv2.resize(img, (self.img_size, self.img_size)) mask cv2.resize(mask, (self.img_size, self.img_size), interpolationcv2.INTER_NEAREST) # 归一化到 [0, 1] 并转成 CHW 格式 img img.astype(np.float32) / 255.0 img img.transpose(2, 0, 1) mask mask.astype(np.float32) / 255.0 mask np.expand_dims(mask, axis0) return torch.from_numpy(img), torch.from_numpy(mask)上面这段代码有几个值得注意的处理细节。第一缩放到 512 时mask 的插值方式必须是INTER_NEAREST不能用双线性插值否则二值掩膜的边界会出现灰阶过渡值后续计算 IoU 时阈值分割会很麻烦。第二归一化只对图像做mask 不做减均值除方差的操作因为它要保持原始的二值语义。第三数据路径用了 glob 加 sort这在 Windows 和 Linux 下行为一致不容易踩路径分隔符的坑。3.2 训练脚本的运行方式与核心参数调整不带任何参数直接跑python train.py理论上是可以运行的因为作者在代码里做了默认配置。但实际训练前我建议把几个关键超参显式确认一遍。# 激活你的虚拟环境后执行 conda activate transunet cd Transformer-Unet # 直接开始训练会自动读取默认配置 python train.py # 如果你想修改 batch size 和 epoch 数在 train.py 中找到对应参数 python train.py --batch_size 4 --epochs 150 --lr 1e-4train.py 内部的训练循环大体上是这样的结构for epoch in range(epochs): model.train() train_loss 0.0 train_iou 0.0 for imgs, masks in train_loader: imgs, masks imgs.to(device), masks.to(device) # 前向传播 - 计算损失 - 反向传播 outputs model(imgs) loss criterion(outputs, masks) optimizer.zero_grad() loss.backward() optimizer.step() train_loss loss.item() train_iou compute_iou(outputs, masks) # 每个 epoch 结束在验证集上跑一次 val_loss, val_iou validate(model, val_loader, criterion) # 记录日志并保存可视化结果 scheduler.step(val_loss) writer.add_scalar(Loss/train, train_loss / len(train_loader), epoch) writer.add_scalar(IoU/val, val_iou, epoch)代码里的compute_iou和validate函数在 utils.py 中有具体实现。这里想提醒的是学习率衰减策略的配置逻辑项目里通常使用 ReduceLROnPlateau它会监控验证集 loss连续几个 epoch 不下降就衰减学习率。实际训练时如果你发现 loss 曲线反复震荡但 val_iou 长时间不动可以先手动把初始学习率调到 5e-5 再跑一轮试试TransUnet 对学习率比纯 CNN 敏感得多。3.3 训练日志与可视化曲线的解读训练跑完以后项目会输出一系列图像和日志文件。看这些产物比看终端输出重要得多因为它们直接反映了模型是否在正常收敛。正常的曲线模式是训练集 loss 稳步下降验证集 loss 先降后微升过拟合信号IoU 曲线逐步上升并趋于平坦。如果你的训练集 loss 在下降但验证集 loss 从一开始就在涨说明模型容量过大而数据量太少这时优先考虑加数据增强和增大 dropout而不是调小模型。另一个常见现象是学习率衰减曲线呈现阶梯状下降。这是 ReduceLROnPlateau 的正常行为不用紧张。整个训练过程建议用 TensorBoard 或直接把曲线图存本地训练完以后隔天再回看往往能发现当时没注意的过拟合苗头。4. 评估与推理的实操方法指标计算和掩膜输出4.1 evaluate.py 能算出哪些指标评估阶段的目标是回答一个问题模型在没见过的测试数据上表现如何。evaluate.py 会计算四类核心指标每一类都有它的工程意义。IoUIntersection over Union是分割任务最重要的指标它计算预测区域和真实区域的交并比。对血管分割来说IoU 超过 0.65 已经算是不错的水平因为血管像素占比本来就很小稍微偏移一点 IoU 就会掉得很难看。Pixel Accuracy 是所有像素中预测正确的比例这个指标在血管分割中参考价值有限因为背景占绝对多数模型只要全部预测为背景就能拿到 90% 以上的准确率。Recommend、Precision 和 F1-score 需要组合来看——Precision 高说明预测的血管像素大部分是真实的Recall 高说明真实血管大部分被找出来了两者往往此消彼长。4.2 用 predict.py 对单张图片做推理并生成掩膜predict.py 的设计目标很明确给定一张眼底图输出血管分割的掩膜和叠加可视化图。# predict.py 的核心推理流程关键部分保留 def predict_single_image(model, img_path, device): # 读取图像并做同步预处理 img cv2.imread(img_path) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) h, w img.shape[:2] # 缩放到模型输入尺寸 img_resized cv2.resize(img, (512, 512)) img_tensor torch.from_numpy(img_resized).float().permute(2, 0, 1).unsqueeze(0) / 255.0 # 推理 model.eval() with torch.no_grad(): output model(img_tensor.to(device)) pred torch.sigmoid(output).squeeze().cpu().numpy() # 二值化和尺寸还原 pred_binary (pred 0.5).astype(np.uint8) * 255 pred_binary cv2.resize(pred_binary, (w, h), interpolationcv2.INTER_NEAREST) # 叠加可视化 overlay img.copy() overlay[pred_binary 255] [0, 255, 0] return pred_binary, overlay推理时最容易犯的错误是训练时做了 resize推理时忘记做尺寸还原。上面代码中用cv2.resize把预测结果恢复到原图尺寸这一步不能省。另外注意阈值的选择代码里默认 0.5如果你发现预测的血管普遍偏细或偏粗可以往下或往上调一调这个阈值在实际项目中经常会根据自己的数据调整到 0.4 到 0.6 之间。4.3 从训练到评估的完整流程串跑从头到尾跑一遍的标准姿势是先用默认参数训练 100 个 epoch 左右然后立刻跑 evaluate.py 拿到 test 指标再用 predict.py 挑两张典型图看可视化效果。三个脚本串起来没有依赖障碍唯一的坑是路径配置——train.py 里指定的输出目录和 evaluate.py 里指定的权重路径必须一致。一个可行的做法是训练完以后把 best_model.pth 单独复制到checkpoints/目录评估和推理统一从这个目录读取。这样即使后续训练了新模型也能随时切回旧权重做对比实验。5. 避坑指南TransUnet 在 DRIVE 上最容易踩的五个坑5.1 坑一直接跑 train.py 报显存不足现象启动训练后几秒钟报CUDA out of memory终端直接跳出错误信息。原因TransUnet 的 Transformer 编码器部分参数量大默认 batch size 可能是 8 或 16在 8GB 显存上直接跑不不动。解决把 batch size 改到 2 或 4同时把图像尺寸从 512 降到 384。如果你不想改代码也可以在训练命令里临时传参前提是 train.py 预留了命令行接口。再不行就在 dataset.py 里加一个img_size384的参数测试一下DRIVE 原图是 565×584缩到 384 信息损失不大但对显存的友好程度提升非常明显。5.2 坑二安装 torch 版本和代码接口不兼容现象在导入项目依赖时终端报出各种module has no attribute比如torch.nn.functional.interpolate新版本行为改变或者某个 API 在高版本 torch 中已废弃。原因项目是在 Python 3.10 环境下打包的torch 版本大概率是 1.12 到 2.x 之间。你本地如果装的是 torch 2.4 或更高部分接口的默认参数行为可能发生了变化。解决建议先按照 requirements.txt 里的版本安装如果 requirements.txt 里没有锁死版本号就按 torch 1.13.1cu117 来装这个版本和 Python 3.10 的兼容性较好且 TransUnet 的前向传播不会遇到接口问题。不要一开始就上最新版 torch。5.3 坑三loss 降低但 IoU 始终上不去现象训练集 loss 正常下降但 val_iou 始终在 0.5 附近波动甚至不涨。原因血管分割的类别极度不均衡血管像素只占整张图的 10% 左右。普通的 BCE Loss 会把大量梯度用在背景像素上模型学到的是全预测为背景的局部最优解。解决损失函数换成 Dice Loss 或者 BCE Dice 的组合损失。Dice Loss 直接优化 IoU 相关的目标对类别不均衡问题更友好。实践中 BCE 0.5*Dice 的组合在血管分割上通常能稳定提升 5 到 10 个百分点的 IoU。不要只盯着 loss 数值loss 降到 0.1 以下时 IoU 可能只有 0.6两者没有严格的线性对应关系。5.4 坑四测试集指标和训练集指标差距巨大现象训练集 IoU 高达 0.8测试集 IoU 只有 0.5 左右明显过拟合。原因DRIVE 数据量不大训练集只有 20 张图左右TransUnet 参数量大很容易把训练集的细节背下来。解决三个手段叠加使用——加大数据增强随机旋转、翻转、弹性形变、提高 dropout 比例到 0.15 到 0.2、使用早停法。其中数据增强的效果最明显。另外可以考虑在 dataset.py 里对训练集做随机裁切增加有效样本数。5.5 坑五推理结果出现棋盘格噪声现象predict.py 输出的掩膜图上血管边缘有规律的棋盘格状伪影。原因这是转置卷积Transposed Convolution叠加产生的重叠效应。TransUnet 解码器里用了多层上采样转置卷积的 stride 和 kernel size 不匹配就会产生这种噪声。解决把解码器的上采样方式从上采样卷积替换为双线性插值卷积或者使用 PixelShuffle。如果不想改模型可以给预测结果加一个轻度的中值滤波但这是治标不治本。从工程角度看微调上采样方式不仅解决棋盘格问题还能轻微收敛边缘精度。6. 进阶玩法把 TransUnet 迁移到自己的数据集并稳定复现DRIVE 跑通只是第一步真正的价值在于把这套流程迁移到自己的数据上。核心操作在 README 里也有提过但有几个细节值得展开。迁移的第一步是数据格式对齐。你的数据集至少要具备两个目录images 存放输入的 RGB 图masks 存放对应的二值掩膜图。图片尺寸不需要和 DRIVE 一致dataset.py 会统一缩放到目标尺寸。文件命名也不需要特殊规则只要 images 和 masks 下的文件用相同的文件名前缀就行。第二步是关注类别是否只有两类。如果你的任务是多类别分割需要改三处model 的 decoder 输出通道改成类别数、dataset.py 的 mask 编码逻辑改成多通道 one-hot、loss 换成 CrossEntropyLoss。TransUnet 的结构天然支持多类别分割但代码里的 DC 指标计算函数需要同步调整。验证方法的建议是训练完成后不要只跑一次测试集就下结论用五折交叉验证跑五次取 IoU 的平均值和标准差。实际工程中这个标准差的参考价值往往比平均值还大——如果你的模型在不同数据划分下 IoU 波动超过 3 个百分点说明泛化能力还不够稳定。从那以后我每次接分割相关的项目都会把这种训练、评估、推理三段式 一次性输出全套可视化曲线的流程先跑一遍不管换什么数据集和模型结构先把管道打通再看效果。希望这份项目拆解对你有用。本文还有配套的精品资源点击获取 TransUnet实战：从DRIVE分割到自定义数据集迁移
简介面向医学图像分割研究者与深度学习开发者这份基于TransUnet的DRIVE视网膜血管分割实战资源包含了完整可运行代码和已标注背景/前景的DRIVE数据集可端到端完成从模型训练到评估推理的全流程。压缩包共76个文件整体仅7.87MB其中40个png文件用于数据集样本与分割结果可视化18个py源码文件覆盖训练、验证、推理三大核心脚本并包含Transformer、Unet、vit等关键模块15个pyc为运行缓存另附requirements.txt、README与txt说明文档借助目录结构和注释可快速定位各功能模块。目前已有323人浏览学习适合医学图像领域入门及进阶者。代码注释非常详细train脚本会生成训练集和验证集的loss、IoU曲线、学习率衰减曲线、训练日志以及数据集可视化图像evaluate脚本可计算测试集的IoU、recall、precision、像素准确率等评估指标predict脚本支持单张图像推理输出GT掩膜与GT叠加原图的融合图便于直观检验分割效果。配套README给出了傻瓜式训练步骤按说明修改路径即可迁移到自定义数据实用性强。1. TransUnet 在 DRIVE 上的分割实战拿到开源项目后到底该怎么跑通、怎么改、怎么评估眼底血管分割是医学影像分割里最经典的入门场景DRIVE 数据集也常被拿来当作算法验证的基准。这个项目给的是一套基于 TransUnet 的完整实现不是只有模型结构而是把训练、评估、推理三个环节全做了闭环。你拿到手里会发现train.py 负责产出训练日志和曲线evaluate.py 负责算测试集指标predict.py 负责出掩膜图三件事分得清清楚楚。更友好的一点是代码里做了详细注释README 里写了训练自定义数据的流程基本能称得上傻瓜式运行。适合三类人刚接触语义分割、想用 Transformer 结构跑通一条完整流程的研究生需要拿 DRIVE 做 baseline 对比的算法工程师以及想快速迁移到自有数据集、但不想从零搭框架的从业者。2. 项目解包文件清单与三个脚本的分工逻辑2.1 压缩包里到底有什么解开压缩包后你会看到两套核心目录data/下存放数据集相关结构和代码unet/下则是模型定义。先把这个结构理清楚比急着跑 train.py 重要得多。Transformer-Unet/ ├── data/ │ ├── train/ # 训练集数据按 DRIVE 原始划分存放 │ ├── val/ # 验证集数据 │ └── test/ # 测试集数据 ├── unet/ │ ├── __init__.py │ ├── vanilla_transformer.py # 原始 Transformer 模块 │ ├── unet_transformer.py # TransUnet 的核心编码器-解码器结构 │ └── common.py # 公共组件如 patch embedding ├── train.py # 训练入口 ├── evaluate.py # 测试集评估入口 ├── predict.py # 推理入口 ├── utils.py # 工具函数 ├── dataset.py # 数据加载与预处理 ├── confuse_matrix.py # 混淆矩阵与指标计算 ├── requirements.txt # 依赖清单 └── README.md # 使用说明文件层级不深但每个文件的责任边界很清楚。第一次接触这个项目时建议先读 README 和 dataset.py再去看模型定义顺序反了容易一头扎进 Transformer 的细节里出不来。2.2 三个入口脚本分别承担什么职责这个项目最大的特点是把训练、评估、推理拆成了三个独立脚本不像很多开源项目把全部逻辑塞进一个 main 函数。这种拆法在实际工作中更实用——训练和评估解耦后你可以随时用同一套权重去验证不同的测试集不需要重新训练。train.py 做的事情不只是迭代数据更新参数它会在训练过程中持续产出可视化和日志文件训练集和验证集的 loss 曲线、IoU 曲线、学习率衰减曲线以及数据集的原始图像和标注可视化图。这意味着训练结束后你不需要额外写代码就能复盘整个训练过程。confuse_matrix.py 是评估阶段的核心依赖。evaluate.py 调用它计算测试集上的 IoU、Recall、Precision、像素准确率等指标最终输出的不是一两个数字而是一套完整的混淆矩阵指标。这些指标对医学影像分割尤其重要因为背景像素占绝对多数只看准确率会严重虚高。predict.py 负责最后的推理。它会读入测试图像输出两个结果gt_mask真实标注掩膜和预测的掩膜同时生成将预测叠加在原图上的可视化图像。这个设计很贴心因为分割任务的效果好不好肉眼看叠加图往往比看数值更直观。2.3 依赖环境与配置要点requirements.txt 里列出了项目运行所需的库。从实际经验看最核心的依赖集中在 PyTorch、TorchVision、NumPy、OpenCV 和 Matplotlib 上。如果你之前跑过其他深度学习项目大概率不需要重新安装太多东西。需要特别注意的是 Python 版本。这个项目里出现了dataset.cpython-310.pyc文件说明作者是在 Python 3.10 环境下打包的。你本地最好也用 3.10 或更高版本避免因为 Python 版本差异导致部分语法或库接口不兼容。CUDA 环境建议用 11.x 以上版本显存至少 8GB因为 TransUnet 的参数量比纯 CNN 结构大不少。3. 跑通训练全流程从数据加载到指标曲线3.1 DRIVE 数据集的加载与预处理逻辑dataset.py 负责把磁盘上的图像读入内存并做预处理。这里最关键的是搞清楚 DRIVE 数据集的格式它包含训练集、验证集和测试集每张图都有对应的血管标注图。DRIVE 的原始图像是 565×584 的 RGB 眼底图标注图是单通道二值图。# dataset.py 中核心的数据加载逻辑简化版 import cv2 import numpy as np import torch from torch.utils.data import Dataset class DriveDataset(Dataset): def __init__(self, img_dir, mask_dir, img_size512, transformNone): self.img_paths sorted(glob.glob(os.path.join(img_dir, *.png))) self.mask_paths sorted(glob.glob(os.path.join(mask_dir, *.png))) self.img_size img_size self.transform transform def __len__(self): return len(self.img_paths) def __getitem__(self, idx): img cv2.imread(self.img_paths[idx]) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) mask cv2.imread(self.mask_paths[idx], cv2.IMREAD_GRAYSCALE) # 统一缩放到模型输入尺寸 img cv2.resize(img, (self.img_size, self.img_size)) mask cv2.resize(mask, (self.img_size, self.img_size), interpolationcv2.INTER_NEAREST) # 归一化到 [0, 1] 并转成 CHW 格式 img img.astype(np.float32) / 255.0 img img.transpose(2, 0, 1) mask mask.astype(np.float32) / 255.0 mask np.expand_dims(mask, axis0) return torch.from_numpy(img), torch.from_numpy(mask)上面这段代码有几个值得注意的处理细节。第一缩放到 512 时mask 的插值方式必须是INTER_NEAREST不能用双线性插值否则二值掩膜的边界会出现灰阶过渡值后续计算 IoU 时阈值分割会很麻烦。第二归一化只对图像做mask 不做减均值除方差的操作因为它要保持原始的二值语义。第三数据路径用了 glob 加 sort这在 Windows 和 Linux 下行为一致不容易踩路径分隔符的坑。3.2 训练脚本的运行方式与核心参数调整不带任何参数直接跑python train.py理论上是可以运行的因为作者在代码里做了默认配置。但实际训练前我建议把几个关键超参显式确认一遍。# 激活你的虚拟环境后执行 conda activate transunet cd Transformer-Unet # 直接开始训练会自动读取默认配置 python train.py # 如果你想修改 batch size 和 epoch 数在 train.py 中找到对应参数 python train.py --batch_size 4 --epochs 150 --lr 1e-4train.py 内部的训练循环大体上是这样的结构for epoch in range(epochs): model.train() train_loss 0.0 train_iou 0.0 for imgs, masks in train_loader: imgs, masks imgs.to(device), masks.to(device) # 前向传播 - 计算损失 - 反向传播 outputs model(imgs) loss criterion(outputs, masks) optimizer.zero_grad() loss.backward() optimizer.step() train_loss loss.item() train_iou compute_iou(outputs, masks) # 每个 epoch 结束在验证集上跑一次 val_loss, val_iou validate(model, val_loader, criterion) # 记录日志并保存可视化结果 scheduler.step(val_loss) writer.add_scalar(Loss/train, train_loss / len(train_loader), epoch) writer.add_scalar(IoU/val, val_iou, epoch)代码里的compute_iou和validate函数在 utils.py 中有具体实现。这里想提醒的是学习率衰减策略的配置逻辑项目里通常使用 ReduceLROnPlateau它会监控验证集 loss连续几个 epoch 不下降就衰减学习率。实际训练时如果你发现 loss 曲线反复震荡但 val_iou 长时间不动可以先手动把初始学习率调到 5e-5 再跑一轮试试TransUnet 对学习率比纯 CNN 敏感得多。3.3 训练日志与可视化曲线的解读训练跑完以后项目会输出一系列图像和日志文件。看这些产物比看终端输出重要得多因为它们直接反映了模型是否在正常收敛。正常的曲线模式是训练集 loss 稳步下降验证集 loss 先降后微升过拟合信号IoU 曲线逐步上升并趋于平坦。如果你的训练集 loss 在下降但验证集 loss 从一开始就在涨说明模型容量过大而数据量太少这时优先考虑加数据增强和增大 dropout而不是调小模型。另一个常见现象是学习率衰减曲线呈现阶梯状下降。这是 ReduceLROnPlateau 的正常行为不用紧张。整个训练过程建议用 TensorBoard 或直接把曲线图存本地训练完以后隔天再回看往往能发现当时没注意的过拟合苗头。4. 评估与推理的实操方法指标计算和掩膜输出4.1 evaluate.py 能算出哪些指标评估阶段的目标是回答一个问题模型在没见过的测试数据上表现如何。evaluate.py 会计算四类核心指标每一类都有它的工程意义。IoUIntersection over Union是分割任务最重要的指标它计算预测区域和真实区域的交并比。对血管分割来说IoU 超过 0.65 已经算是不错的水平因为血管像素占比本来就很小稍微偏移一点 IoU 就会掉得很难看。Pixel Accuracy 是所有像素中预测正确的比例这个指标在血管分割中参考价值有限因为背景占绝对多数模型只要全部预测为背景就能拿到 90% 以上的准确率。Recommend、Precision 和 F1-score 需要组合来看——Precision 高说明预测的血管像素大部分是真实的Recall 高说明真实血管大部分被找出来了两者往往此消彼长。4.2 用 predict.py 对单张图片做推理并生成掩膜predict.py 的设计目标很明确给定一张眼底图输出血管分割的掩膜和叠加可视化图。# predict.py 的核心推理流程关键部分保留 def predict_single_image(model, img_path, device): # 读取图像并做同步预处理 img cv2.imread(img_path) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) h, w img.shape[:2] # 缩放到模型输入尺寸 img_resized cv2.resize(img, (512, 512)) img_tensor torch.from_numpy(img_resized).float().permute(2, 0, 1).unsqueeze(0) / 255.0 # 推理 model.eval() with torch.no_grad(): output model(img_tensor.to(device)) pred torch.sigmoid(output).squeeze().cpu().numpy() # 二值化和尺寸还原 pred_binary (pred 0.5).astype(np.uint8) * 255 pred_binary cv2.resize(pred_binary, (w, h), interpolationcv2.INTER_NEAREST) # 叠加可视化 overlay img.copy() overlay[pred_binary 255] [0, 255, 0] return pred_binary, overlay推理时最容易犯的错误是训练时做了 resize推理时忘记做尺寸还原。上面代码中用cv2.resize把预测结果恢复到原图尺寸这一步不能省。另外注意阈值的选择代码里默认 0.5如果你发现预测的血管普遍偏细或偏粗可以往下或往上调一调这个阈值在实际项目中经常会根据自己的数据调整到 0.4 到 0.6 之间。4.3 从训练到评估的完整流程串跑从头到尾跑一遍的标准姿势是先用默认参数训练 100 个 epoch 左右然后立刻跑 evaluate.py 拿到 test 指标再用 predict.py 挑两张典型图看可视化效果。三个脚本串起来没有依赖障碍唯一的坑是路径配置——train.py 里指定的输出目录和 evaluate.py 里指定的权重路径必须一致。一个可行的做法是训练完以后把 best_model.pth 单独复制到checkpoints/目录评估和推理统一从这个目录读取。这样即使后续训练了新模型也能随时切回旧权重做对比实验。5. 避坑指南TransUnet 在 DRIVE 上最容易踩的五个坑5.1 坑一直接跑 train.py 报显存不足现象启动训练后几秒钟报CUDA out of memory终端直接跳出错误信息。原因TransUnet 的 Transformer 编码器部分参数量大默认 batch size 可能是 8 或 16在 8GB 显存上直接跑不不动。解决把 batch size 改到 2 或 4同时把图像尺寸从 512 降到 384。如果你不想改代码也可以在训练命令里临时传参前提是 train.py 预留了命令行接口。再不行就在 dataset.py 里加一个img_size384的参数测试一下DRIVE 原图是 565×584缩到 384 信息损失不大但对显存的友好程度提升非常明显。5.2 坑二安装 torch 版本和代码接口不兼容现象在导入项目依赖时终端报出各种module has no attribute比如torch.nn.functional.interpolate新版本行为改变或者某个 API 在高版本 torch 中已废弃。原因项目是在 Python 3.10 环境下打包的torch 版本大概率是 1.12 到 2.x 之间。你本地如果装的是 torch 2.4 或更高部分接口的默认参数行为可能发生了变化。解决建议先按照 requirements.txt 里的版本安装如果 requirements.txt 里没有锁死版本号就按 torch 1.13.1cu117 来装这个版本和 Python 3.10 的兼容性较好且 TransUnet 的前向传播不会遇到接口问题。不要一开始就上最新版 torch。5.3 坑三loss 降低但 IoU 始终上不去现象训练集 loss 正常下降但 val_iou 始终在 0.5 附近波动甚至不涨。原因血管分割的类别极度不均衡血管像素只占整张图的 10% 左右。普通的 BCE Loss 会把大量梯度用在背景像素上模型学到的是全预测为背景的局部最优解。解决损失函数换成 Dice Loss 或者 BCE Dice 的组合损失。Dice Loss 直接优化 IoU 相关的目标对类别不均衡问题更友好。实践中 BCE 0.5*Dice 的组合在血管分割上通常能稳定提升 5 到 10 个百分点的 IoU。不要只盯着 loss 数值loss 降到 0.1 以下时 IoU 可能只有 0.6两者没有严格的线性对应关系。5.4 坑四测试集指标和训练集指标差距巨大现象训练集 IoU 高达 0.8测试集 IoU 只有 0.5 左右明显过拟合。原因DRIVE 数据量不大训练集只有 20 张图左右TransUnet 参数量大很容易把训练集的细节背下来。解决三个手段叠加使用——加大数据增强随机旋转、翻转、弹性形变、提高 dropout 比例到 0.15 到 0.2、使用早停法。其中数据增强的效果最明显。另外可以考虑在 dataset.py 里对训练集做随机裁切增加有效样本数。5.5 坑五推理结果出现棋盘格噪声现象predict.py 输出的掩膜图上血管边缘有规律的棋盘格状伪影。原因这是转置卷积Transposed Convolution叠加产生的重叠效应。TransUnet 解码器里用了多层上采样转置卷积的 stride 和 kernel size 不匹配就会产生这种噪声。解决把解码器的上采样方式从上采样卷积替换为双线性插值卷积或者使用 PixelShuffle。如果不想改模型可以给预测结果加一个轻度的中值滤波但这是治标不治本。从工程角度看微调上采样方式不仅解决棋盘格问题还能轻微收敛边缘精度。6. 进阶玩法把 TransUnet 迁移到自己的数据集并稳定复现DRIVE 跑通只是第一步真正的价值在于把这套流程迁移到自己的数据上。核心操作在 README 里也有提过但有几个细节值得展开。迁移的第一步是数据格式对齐。你的数据集至少要具备两个目录images 存放输入的 RGB 图masks 存放对应的二值掩膜图。图片尺寸不需要和 DRIVE 一致dataset.py 会统一缩放到目标尺寸。文件命名也不需要特殊规则只要 images 和 masks 下的文件用相同的文件名前缀就行。第二步是关注类别是否只有两类。如果你的任务是多类别分割需要改三处model 的 decoder 输出通道改成类别数、dataset.py 的 mask 编码逻辑改成多通道 one-hot、loss 换成 CrossEntropyLoss。TransUnet 的结构天然支持多类别分割但代码里的 DC 指标计算函数需要同步调整。验证方法的建议是训练完成后不要只跑一次测试集就下结论用五折交叉验证跑五次取 IoU 的平均值和标准差。实际工程中这个标准差的参考价值往往比平均值还大——如果你的模型在不同数据划分下 IoU 波动超过 3 个百分点说明泛化能力还不够稳定。从那以后我每次接分割相关的项目都会把这种训练、评估、推理三段式 一次性输出全套可视化曲线的流程先跑一遍不管换什么数据集和模型结构先把管道打通再看效果。希望这份项目拆解对你有用。本文还有配套的精品资源点击获取