TransUnet这个网络常跑医学图像分割的朋友应该都不陌生。它把CNN的特征提取能力和Transformer的全局建模能力拼在一起在不少分割任务上都拿到了不错的效果现在很多论文还是会拿它当对比基准。但这里有个很现实的问题官方代码默认是跑Synapse、ACDC这类医学灰度影像的数据加载部分直接读的是npy或者h5格式预处理链路也完全围绕灰度图设计。如果你手里只有一堆彩色RGB图片直接拿官方train.py开跑大概率会在数据读取阶段就报错或者训练出来的模型根本学不到东西。这篇文章把我自己的改造过程完整拆一遍从数据集整理、官方代码结构解读、RGB三通道适配到训练参数调整和推理落地一条龙讲清楚帮想用TransUnet跑自己RGB分割数据的朋友少走弯路。1. 先弄清TransUnet做了什么以及你的RGB数据为什么不能直接喂给它1.1 网络结构核心思路TransUnet的结构可以简单拆成三块CNN编码器、Transformer编码器、U-Net解码器。CNN部分通常用ResNet50作为骨干先对输入图像做下采样提取局部特征随后把CNN输出的特征图切成固定大小的patch展平成序列后送入Transformer的Encoder里用自注意力机制去捕获像素之间的长距离依赖关系最后再把Transformer输出的序列重新还原成特征图的形状交给U-Net风格的解码器逐步上采样恢复出和原图分辨率一致的分割结果。这种混合设计的核心理由其实很直白纯Transformer擅长建模全局关系但对像素级细节有点粗枝大叶直接用在分割上容易出现边界糊成一片的问题纯CNN又受限于感受野很难把相隔很远的同类区域连起来。TransUnet相当于用CNN先打底、用Transformer补全局视野再用U-Net把细节捞回来三个环节各司其职。实际训练中被称作“三明治”式的混合编码器带来的效果改善在医学影像这种结构复杂、边界模糊的任务上尤其明显。1.2 官方代码默认的数据流与RGB数据的冲突点TransUnet官方仓库的数据处理流程核心是围绕Synapse多器官分割数据集设计的。官方数据加载模块会把原始影像和标签都转成npy格式的数组然后按2D切片的方式一张张喂给网络。这个流程里默认的假设是输入是灰度图数据增强、归一化、通道维度的处理全按单通道来。而Synapse这类数据集标签图里的每个像素值对应一个器官类别编号整体是索引标注图而不是彩色标注图。RGB三通道彩色图替换进来之后主要会撞上三堵墙数据加载模块读图方式不对。官方代码很多地方直接假设输入已经是npy数组而不是从JPG、PNG这类常见图片格式里读RGB值。归一化参数不对。灰度图的归一化一般只针对单通道做而RGB图像需要逐通道做标准化否则模型输入分布完全偏离预训练权重的分布收敛速度会慢很多。标签处理逻辑不对。很多自己标的数据集标注文件保存成彩色PNG每个类别用一种颜色这类彩色标注必须转成0、1、2...这样的类别索引图才能配合交叉熵这类损失函数使用。所以这里有一个很关键的结论TransUnet的网络结构本身是能吃RGB三通道输入的ViT部分也有对应的3通道patch embedding真正需要改的是官方代码里“从文件到网络输入”之间的这一整条数据链路。这篇文章后面讲的所有改动本质上都是围绕这条链路来做手术。2. 训练环境准备与数据集整理规范2.1 硬件环境和依赖安装我自己用的是一张RTX 409024GB显存batch size可以开到16配合224x224输入尺寸跑150个epoch大约需要六到八个小时。如果你手里的显卡是12GB显存建议batch size先降到8或者用梯度累积来模拟更大的batch。CPU内存方面32GB比较稳妥因为数据增强阶段会在内存里同时解压多张图处理遇到大佬数据集动辄上万张图时内存小了很容易被系统杀掉进程。环境依赖上官方代码仓库要求的核心依赖是python 3.8 pytorch 1.10 torchvision timm einops tensorboard SimpleITK跑Synapse原版数据时需要跑纯RGB数据可暂时不装装完基础依赖后我额外建议安装opencv-python和albumentations前者读取图片方便、速度快后者做图像增强时能保持image和mask同步变换比手动写增强函数省心很多。安装命令如下pip install opencv-python albumentations tensorboard einops timm2.2 数据集目录结构设计整理自己的RGB分割数据集我强烈建议一开始就按下面的目录规范来组织后面改代码、写加载器、做评估都会省很多事datasets/ ├── RGBDataset/ │ ├── train/ │ │ ├── images/ # 训练原图jpg/png均可3通道 │ │ │ ├── img_001.jpg │ │ │ └── img_002.jpg │ │ └── masks/ # 训练标签PNG格式单通道索引图或彩色标注图 │ │ ├── img_001.png │ │ └── img_002.png │ ├── val/ │ │ ├── images/ │ │ └── masks/ │ ├── test/ │ │ ├── images/ │ │ └── masks/ │ └── train.txt # 每一行图片路径 制表符 标签路径 └── pretrained/ ├── R50ViT-B_16.npz # TransUnet官方提供的ImageNet预训练权重 └── R50ViT-B_16.json关于标签的格式这里有个重点需要单独拎出来说。我见过太多人在这个环节掉坑用Labelme或者PS标注完之后保存出来的mask是彩色图比如类别0是黑色、类别1是红色、类别2是绿色每个类别对应一种RGB向量。这种情况下不能直接把maks图喂给网络因为网络输出的类别索引是0、1、2这种一维数字和RGB颜色向量对不上。必须先把彩色标注图转成单通道索引图做法是把每个RGB值映射到一个类id上。如果标注图是单通道PNG里面每个像素的灰度值恰好等于类别id那就省事了直接读就行。在Python里用opencv读图时注意一点cv2.imread(path, cv2.IMREAD_GRAYSCALE)读出来的是单通道灰度图cv2.imread(path)默认读出来的是BGR三通道彩色图这两个模式千万别搞混。2.3 数据划分与类别映射表划分数据集时我建议按**训练集70%、验证集15%、测试集15%**的比例随机划分。随机划分前先检查一下每张图的类别分布确保验证集和测试集里每个类别都存在不然评估阶段算出来的mIoU / Dice会忽高忽低没有参考意义。更精细一点的做法是按图像来源分组划分比如某一批图来自同一台设备或者同一时段拍摄就把它们放到同一组里再按组切分防止数据泄漏导致的指标虚高。手动标注出来的RGB分割数据集类别数量从二分类到十来个类别都很常见。通常需要维护一张类别映射表类似这样类别id像素颜色RGB含义0(0, 0, 0)背景1(255, 0, 0)目标物A2(0, 255, 0)目标物B3(0, 0, 255)目标物C写转换脚本时优先把所有像素RGB值用numpy的矩阵运算一次性映射到类别id不要用Python循环逐像素判断后者在1920x1080这种分辨率下会慢到怀疑人生。代码示例如下import numpy as np import cv2 color_to_id { (0, 0, 0): 0, (255, 0, 0): 1, (0, 255, 0): 2, (0, 0, 255): 3, } def rgb_mask_to_index(mask_path, out_path): # 读取彩色标注图注意OpenCV默认是BGR顺序 mask_bgr cv2.imread(mask_path) mask_rgb cv2.cvtColor(mask_bgr, cv2.COLOR_BGR2RGB) h, w mask_rgb.shape[:2] # 初始化索引图-1表示未知类别方便后面检查漏标 index_map np.full((h, w), fill_value-1, dtypenp.int32) for cls_id, rgb in color_to_id.items(): match np.all(mask_rgb np.array(rgb).reshape(1, 1, 3), axis-1) index_map[match] cls_id if np.any(index_map -1): print(f警告: {mask_path} 中存在未归类像素) # 保存为单通道PNG cv2.imwrite(out_path, index_map.astype(np.uint8))3. 核心代码改造把官方数据链路换成RGB版本3.1 官方代码结构导读改代码之前先把官方仓库的文件布局摸清楚。TransUnet官方仓库里与训练直接相关的关键文件大概有这几个train.py # 训练入口 test.py # 测试入口 datasets/ ├── dataset_synapse.py # Synapse数据集加载器 └── dataset_acdc.py # ACDC数据集加载器 lib/models/transunet.py # TransUnet模型定义 utils/ ├── losses.py # 损失函数 ├── metrics.py # 评估指标 └── test.py # 测试工具函数 lists/ ├── lists_Synapse/ │ ├── train.txt │ └── test_vol.txt其中dataset_synapse.py是我花时间最多的地方因为官方代码里大量使用了np.load直接读npy数组并内置了“把3D体数据切成2D切片”的逻辑。我们自己跑RGB图片最省事的方案其实不是去官方代码上修修补补而是直接写一个新的Dataset类把官方数据集类和文件路径都替换掉。3.2 自定义RGB Dataset类下面这份代码是我在实际项目中整理出来的可以直接放到datasets/目录下新建一个dataset_rgb.py文件里import os import cv2 import numpy as np import torch from torch.utils.data import Dataset import albumentations as A from albumentations.pytorch import ToTensorV2 class RGBSegmentationDataset(Dataset): def __init__(self, data_dir, split_file, img_size224, modetrain, num_classes2): self.data_dir data_dir self.mode mode self.img_size img_size self.num_classes num_classes # 读取split文件每行是images/xxx.jpg\tmasks/xxx.png self.samples [] with open(split_file, r, encodingutf-8) as f: for line in f: line line.strip() if not line: continue img_rel, mask_rel line.split(\t) img_path os.path.join(data_dir, img_rel) mask_path os.path.join(data_dir, mask_rel) if os.path.exists(img_path) and os.path.exists(mask_path): self.samples.append((img_path, mask_path)) # 数据增强管线 if mode train: self.transform A.Compose([ A.RandomResizedCrop(heightimg_size, widthimg_size, scale(0.7, 1.0), p0.5), A.HorizontalFlip(p0.5), A.VerticalFlip(p0.2), A.RandomBrightnessContrast(p0.3), A.Normalize( mean[123.675, 116.28, 103.53], std[58.395, 57.12, 57.375], max_pixel_value255.0 ), ToTensorV2() ]) else: self.transform A.Compose([ A.Resize(heightimg_size, widthimg_size), A.Normalize( mean[123.675, 116.28, 103.53], std[58.395, 57.12, 57.375], max_pixel_value255.0 ), ToTensorV2() ]) def __len__(self): return len(self.samples) def __getitem__(self, idx): img_path, mask_path self.samples[idx] # 读取RGB原图cv2默认读成BGR先转回RGB image_bgr cv2.imread(img_path) if image_bgr is None: raise ValueError(f无法读取图像: {img_path}) image_rgb cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB) # 读取标签如果标签是单通道索引图用IMREAD_GRAYSCALE mask cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) if mask is None: raise ValueError(f无法读取标签: {mask_path}) # 标签里如果带有未知类别(255)需要先处理掉 mask np.clip(mask, 0, self.num_classes - 1) # 数据增强image和mask必须保持相同的空间变换 augmented self.transform(imageimage_rgb, maskmask) image_tensor augmented[image] mask_tensor torch.from_numpy(augmented[mask]).long() return image_tensor, mask_tensor这段代码里有几个地方是专门为RGB三通道数据设计的。归一化用的mean和std就是ImageNet预训练权重对应的RGB通道数值顺序是R、G、B。如果你用的是自己从零训练的模型也可以改用自己数据集上统计的均值和方差但大多数情况下用ImageNet的数值已经足够。此外我用albumentations而不是PyTorch官方的torchvision.transforms做增强原因是这个库自带了一套image和mask同步变换的封装。拿随机裁剪来说如果只用torchvision自带的RandomCrop原图和标签如果不用同一个随机种子去裁剪位置会对不上albumentations里直接传两个参数进去它内部帮你处理好同步问题代码简洁、也不容易出bug。这一点在分割任务里非常实用。3.3 修改训练入口文件官方train.py里与数据加载和模型输出相关的几个关键改法我逐个说明。第一个是实例化Dataset和DataLoader的部分。官方代码里创建dataloader的语句可能长这样db_train Dataset(base_dirargs.root_path, list_dirargs.list_dir, splittrain) trainloader torch.utils.data.DataLoader(db_train, batch_sizeargs.batch_size, shuffleTrue, num_workers8, pin_memoryTrue)这行代码里的Dataset指的是官方自带的Synapse_dataset。我们要替换成自己写的RGBSegmentationDataset并且传入split文件路径而不是目录from datasets.dataset_rgb import RGBSegmentationDataset train_dataset RGBSegmentationDataset( data_dirargs.root_path, split_fileargs.list_dir, # 这里改成train.txt文件路径 img_sizeargs.img_size, modetrain, num_classesargs.num_classes ) trainloader torch.utils.data.DataLoader( train_dataset, batch_sizeargs.batch_size, shuffleTrue, num_workers8, pin_memoryTrue, drop_lastTrue )这里drop_lastTrue是个小细节。如果训练集样本数不能被batch_size整除丢弃最后一个不完整的batch能避免BatchNorm层在训练和推理时行为不一致的问题。样本总数不足时这个设置尤其重要。第二个是num_classes参数要传对。官方代码里经常在训练脚本顶部硬编码了输出类别数比如n_classes9对应Synapse的8个器官加背景。改成自己的数据集时务必把这里替换成自己数据集的类别总数包括背景类。否则模型最后一层输出的通道数和标签中的最大类别索引对不上训练时损失直接爆掉报错信息多是“Target size ! Tensor size”之类。第三个是模型定义部分。TransUnet模型在构造时通常需要传入img_size和in_chans等参数net TransUnet( img_size224, in_chans3, num_classesargs.num_classes, embed_dim768, depth12, num_heads12, ... )in_chans3是RGB三通道的关键千万别改成1。虽然官方仓库的预训练模型通常是在3通道ImageNet上预训练的但这个参数还是要显式写清楚。如果你的环境里加载模型时因为timm库版本问题报错可以考虑升级或降级timm到0.4.12左右官方代码依赖的是比较老的版本。3.4 预训练权重的加载处理TransUnet通常会在Vit部分使用ImageNet上预训练的R50ViT-B_16权重。加载预训练权重时有一个容易翻车的点官方提供的R50ViT-B_16.npz文件解压出来的权重dict里键名和模型当前状态的键名可能不一致直接load_state_dict会报“Missing key(s) and unexpected key(s)”的警告。我的习惯是写一个简单的兼容加载函数import numpy as np import torch def load_pretrained_npz(model, npz_path): npz np.load(npz_path, allow_pickleFalse) weights {k: torch.from_numpy(v) for k, v in npz.items()} model_dict model.state_dict() matched {} for k, v in weights.items(): if k in model_dict and model_dict[k].shape v.shape: matched[k] v model_dict.update(matched) model.load_state_dict(model_dict) print(f成功加载预训练权重: {len(matched)} / {len(model_dict)} 层匹配) return model这种加载策略允许网络部分结构不匹配时仍然能加载成功不匹配的层保留随机初始化后面训练时这些层会自己学出来。实测下来用预训练权重做初始化比从零开始训练在收敛速度和最终Dice指标上普遍要好5到10个百分点所以有条件的话尽量把这个权重用上。4. 训练参数配置与训练流程实操4.1 优化器、学习率与损失函数的选择官方代码默认用的是SGD优化器加CosineAnnealing学习率调度。我自己在RGB分割任务上更推荐AdamW CosineAnnealing 线性warmup的组合。SGD在医学分割这种小数据集上收敛偏慢且对学习率敏感AdamW对learning rate没那么挑剔配合warmup预热可以让Transformer部分在训练初期更稳定。我常用的一组基础参数如下优化器: AdamW 初始学习率: 1e-4 权重衰减: 1e-4 batch size: 1624GB显存/ 812GB显存 训练轮数: 150 学习率调度: CosineAnnealing50轮warmup结束到1e-3再降回1e-5如果你执意用SGD学习率可以设置为0.01到0.05之间配合poly学习率策略效果也不错。不过从我测试来看AdamW在分类不平衡的数据集上表现更稳定尤其当你的数据里背景像素占比很大的时候。损失函数部分官方提供了DiceLoss和CrossEntropyLoss两个选择我强烈建议两者结合起来一起用ce_loss nn.CrossEntropyLoss() dice_loss DiceLoss(num_classesargs.num_classes) total_loss ce_loss(logits, labels) dice_loss(logits, labels)单独用CrossEntropyLoss时如果背景像素占90%以上网络很容易把所有像素都预测成背景Dice会很难看单独用DiceLoss时又容易出现loss震荡不收敛的情况。两个loss加起来CE负责提供稳定的梯度方向Dice负责对齐类间平衡配合起来效果最稳。权重比例我一般设1:1如果你的类别极度不均衡可以加大DiceLoss的权重比如0.6的系数给dice、0.4给CE。4.2 训练过程中的监控指标解读训练时我习惯同时打印这几个指标train loss、train dice、val dice、val mIoU。loss能反映模型是否在收敛dice能反映分割质量和背景类占比是否合理。这里说一个经常遇到的“假收敛”现象如果你的分类问题里背景像素占绝大多数模型很快就能学会全部预测成背景这时候train loss看起来在下降但val dice可能只有0.3甚至更低。应对办法就是上面说的DiceLossCrossEntropy组合以及在训练中打印按类别分别统计的dice看清楚模型到底是哪个类别学不动。我在训练过程中的tensorboard日志通常会记录以下几类曲线train/loss train/dice val/dice val/iou val/loss lr其中val/iou里的IoU也就是交并比是衡量分割结果和真实标注重合度的最直观指标。计算方法是每个类别分别算预测正确的像素数 / (预测为该类的像素数 真实为该类的像素数 - 预测正确的像素数)最后对所有类别取平均得到mIoU。如果某个类别在训练集里只有几十个像素它的IoU对整体均值影响不大但你一旦做测试集评估这个类别几乎必然预测失败因此要在训练之前就想好要不要对这样的rare class做上采样或者类别加权。4.3 显存不够时的节流三板斧实测中即使24GB显存如果把分辨率调到512x512batch size基本只能开到4到6。如果显存继续爆掉按下面顺序检查第一把batch size调到2甚至1试试这是最简单直接的办法。batch size小会让BatchNorm的统计量不稳如果你的batch size小于8建议把网络里的BatchNorm换成GroupNorm或者InstanceNorm不然验证时使用running mean指标会有明显跳动。第二开混合精度训练。PyTorch的自动混合精度AMP可以把显存占用砍掉将近一半而且对最终精度基本没有影响。核心改动只有几行代码scaler torch.cuda.amp.GradScaler() for batch_idx, (images, labels) in enumerate(trainloader): images, labels images.cuda(), labels.cuda() optimizer.zero_grad() with torch.cuda.amp.autocast(): logits net(images) loss ce_loss(logits, labels) dice_loss(logits, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()第三把num_workers调低一些。有时候爆显存不是显存真的不够而是DataLoader线程数太多导致内存线程切换开销大num_workers4到8基本够用了不用盲目开满。4.4 训练中断恢复与checkpoint管理训练跑到一半因为服务器重启、显存被占用等原因挂掉是家常便饭。我的习惯是每5个epoch保存一个checkpoint文件名里带上epoch号和val dice并且单独保存一份最新的“last.pth”方便随时从断点恢复。torch.save({ epoch: epoch, model_state_dict: net.state_dict(), optimizer_state_dict: optimizer.state_dict(), scheduler_state_dict: scheduler.state_dict(), best_dice: best_dice, }, checkpoints/last.pth)恢复训练时只需要重新加载last.pth把网络、优化器、调度器的状态全部恢复再继续跑就行。如果数据增强管线里用了random seed记得把np.random.seed和torch.manual_seed也一并固定否则恢复之后虽然模型状态一样但增强出来的数据完全不同训练曲线可能出现小跳变。5. 验证与推理从训练好的模型到一张张预测图5.1 测试指标的计算与可视化训练结束后在测试集上计算Dice、IoU这些指标时要特别注意一点TransUnet输出的是一个形状为B, num_classes, H, W的概率图需要对每个像素在类别维度上做argmax得到最终的类别索引图。然后拿这个索引图和真实标签逐像素比较。def calculate_metrics(pred_mask, true_mask, num_classes): iou_list [] dice_list [] for cls in range(num_classes): pred_cls (pred_mask cls) true_cls (true_mask cls) intersection (pred_cls true_cls).sum() union (pred_cls | true_cls).sum() if union 0: iou_list.append(float(nan)) # 该类别在GT中不存在 else: iou_list.append(intersection / union if union 0 else 0.0) dice_list.append(2 * intersection / (pred_cls.sum() true_cls.sum()) if (pred_cls.sum() true_cls.sum()) 0 else 0.0) return np.nanmean(iou_list), np.nanmean(dice_list)遇到某个类别在真实标签中完全不存在的测试图iou和dice会碰到除零问题。直接跳过该类别还是给0分取决于你评估的目的。如果是给论文写实验对比我一般建议按“该类在ground truth中出现才参与计算”的规则处理并在题注里注明如果是实际产品验证则更倾向于给0分因为漏检的情况同样需要被惩罚。5.2 单张图片推理演示实际部署时经常需要对任意一张新图做在线预测。下面这份脚本可以直接在命令行里跑加载训练好的权重输入一张RGB图片输出预测mask和彩色叠加图import torch import cv2 import numpy as np from lib.models.transunet import TransUnet def predict_image(model, image_path, img_size224, devicecuda): model.eval() image_bgr cv2.imread(image_path) image_rgb cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB) original_h, original_w image_rgb.shape[:2] # resize到网络输入尺寸 image_resized cv2.resize(image_rgb, (img_size, img_size), interpolationcv2.INTER_LINEAR) # 归一化 mean np.array([123.675, 116.28, 103.53], dtypenp.float32) std np.array([58.395, 57.12, 57.375], dtypenp.float32) image_norm (image_resized.astype(np.float32) - mean) / std image_tensor torch.from_numpy(image_norm.transpose(2, 0, 1)).unsqueeze(0).float().to(device) with torch.no_grad(): logits model(image_tensor) # (1, num_classes, img_size, img_size) pred torch.argmax(logits, dim1).squeeze(0).cpu().numpy() # (img_size, img_size) 类别索引 # 恢复到原图分辨率 pred_resized cv2.resize(pred.astype(np.uint8), (original_w, original_h), interpolationcv2.INTER_NEAREST) return pred_resized def save_overlay(image_path, pred_mask, output_path): image_bgr cv2.imread(image_path) overlay image_bgr.copy() # 这里定义简单颜色映射可根据自己的类别调整 color_map { 1: (0, 0, 255), # 红目标A 2: (0, 255, 0), # 绿目标B 3: (255, 0, 0), # 蓝目标C } for cls, color in color_map.items(): overlay[pred_mask cls] color blended cv2.addWeighted(image_bgr, 0.5, overlay, 0.5, 0) cv2.imwrite(output_path, blended)这个脚本里有个关键细节resize预测结果时插值方式必须用INTER_NEAREST不能使用INTER_LINEAR或INTER_CUBIC。因为分割mask是离散的类别编号用线性插值会在类别边界产生不存在的中间值比如类别1和类别2之间出现0.6这种小数值转回uint8后会成为另一个错误类别导致边缘出现一圈鬼影。同理保存在推理阶段所有涉及预测mask的空间变换都要用最近邻插值。5.3 预测结果的后处理与常见问题做完预测后你大概率会看到两类问题。一类是细小孔洞就是预测结果里本该连续的大块区域中间出现零星的小洞。这种情况推荐用简单的形态学闭运算填掉常用的是cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)kernel大小设为3x3或5x5就够太大会把细长结构也给抹掉。另一类是类别边界毛刺感很强。这是像素级分割的通病可以用CRF后处理技术改善但引入CRF会显著增加推理时间实际工程里要权衡。如果你的场景只在乎分割的大致区域对边缘精细度要求不高那就保留原始输出即可不用额外后处理。6. 常见问题与排错实战记录6.1 读图报错与通道顺序错误我自己在最初跑RGB数据时踩的第一个坑就是OpenCV的通道顺序。OpenCV的cv2.imread返回的数组是BGR顺序而PyTorch模型训练时一般期望RGB顺序。如果不做cv2.COLOR_BGR2RGB转换就相当于把红通道和蓝通道对调后送进网络模型训练时loss可能照样下降但验证效果奇差无比因为模型学到的颜色语义完全是错位的。排查方法很简单把读出来的图存一张可视化对比一下看看红色物体是否显示成了蓝色。6.2 标签类别索引与输出类别数不匹配如果标签图中最大像素值是4但num_classes设置成了3训练时会在计算损失时报错Target 4 is out of bounds。相反如果num_classes设置大了而标签里最大类别只有2倒是能跑但模型会浪费大量参数去学一个不存在的类别。所以训练前务必要检查一下标签的最大值mask cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) print(label min:, mask.min(), label max:, mask.max())一旦发现max值等于255几乎可以确定标签是彩色标注图被当成灰度图读进来了需要先做颜色到类别的映射。6.3 验证集Dice高但测试集Dice低这种情况大多不是过拟合而是训练集和测试集分布不一致。RGB分割里最常见的原因是图像的光照条件不同比如训练数据是从晴天场景采集的测试集里全是阴天或者夜晚图像。此时仅靠数据增强是不够的需要针对性做色彩增强比如在训练管线中加入随机色调扰动、灰度化、对比度抖动等。一个更隐蔽的问题是图像分辨率不一致。如果训练时统一resize到了224x224而测试图像是1920x1080再resize回224x224小目标细节信息会损失严重。我的建议是训练时就模拟测试阶段的分辨率策略如果测试时会把整张大图切块预测再拼回来那么训练时也应该用相同尺寸的切块去做增强和裁剪。6.4 训练loss不下降或直接变NaN训练loss变成NaN最常见的原因是学习率设置过大梯度更新一步就越过了数值稳定边界。这时候把学习率降到1e-5到1e-4量级再试情况基本能缓解。第二个常见原因是标签中有超出num_classes范围的异常值。第三个原因是从npy或h5文件里读到了包含NaN的输入数据比如原图某些像素出现坏值训练时梯度回传就会出问题。如果loss一直不下降先别急着调参把训练数据可视化一遍。我之前遇到过数据文件夹里混入了大量损坏的图片OpenCV读取时返回None但代码里没有做检查直接硬算导致loss波动非常离谱。跑训练之前批量扫描一遍图片能否正常读取是个很好的习惯import cv2, os from tqdm import tqdm bad_files [] for root, dirs, files in os.walk(datasets/RGBDataset): for f in files: if f.lower().endswith((.jpg, .png, .jpeg)): path os.path.join(root, f) img cv2.imread(path) if img is None: bad_files.append(path) print(f损坏图片数量: {len(bad_files)})6.5 数据加载慢、GPU利用率上不去训练时GPU利用率只有30%左右CPU反而跑满这是典型的数据加载瓶颈。虽然我们把num_workers开到8了但如果每个worker处理单张图时都要做大量CPU操作比如大尺寸resize、彩色mask转换、复杂增强整体速度依然会被拖慢。几个行之有效的优化手段图片尺寸在进入DataLoader前先统一缩放避免每个worker都在做接近原始分辨率的resize操作。用pin_memoryTrue让GPU直接访问锁页内存减少数据拷贝时间。数据增强管线里避免使用逐像素的Python循环尽量用OpenCV和numpy的向量化操作。如果训练集规模本来就几万张可以考虑在离线阶段把所有图预处理并缓存成LMDB或TFRecord格式训练时直接读取缓存。这个方案改动较大一般数据量小的时候不需要上。6.6 跑通了但分割效果一言难尽怎么办如果训练正常跑完预测时mask却是一团噪点或者大片漏分割大概率不是代码问题而是模型能力没有充分释放。我建议按固定顺序排查先确认评估指标本身可靠。打印出测试集前10张图的预测mask转成灰度图可视化人眼判断预测结果是否至少有大致形状。如果只是边缘粗糙说明模型学到了结构信息可以靠后处理或者继续训练提升。再核查数据集内标注质量。有些标注的边界很粗糙类别和原图对不齐这会导致Dice上限被拉低再怎么调参也上不去。把原图和标注半透明叠加显示检查几处边界是否贴合得干净。最后再考虑模型本身的变化比如增大输入分辨率、增加训练轮数、换更强的预训练权重。TransUnet在224x224输入下对细长结构的分割天然吃亏因为Transformer的patch尺寸是16x16一条只有几个像素宽的线在patch内可能直接糊化。有条件的话可以把输入改成384x384或512x512。显存不够时就用上面的混合精度和梯度累积方案腾空间。7. 写在最后的几点实操建议如果只让我给一个最核心的忠告那就是第一次跑通全流程永远先用最小样本集。我当时先用20张图、10个epoch跑了一遍完整流程确认loss能下降、checkpoint能保存、推理脚本能出图才动用全量数据开正式训练。这一步看着多花了半小时实际上帮你节省的是排查“训练跑了一半才发现mask读错了”这种灾难现场的时间。另外TransUnet官方代码本身是研究性质的代码训练脚本写得相对粗糙很多地方没有做错误处理。改造时不要有“官方的一定是完美”的执念按自己数据的特点去改数据加载器和训练策略才是正路。这篇文章里给的Dataset类和推理脚本是我在RGB分割任务上跑过多次的基础版本你可以直接拿去改路径跑通再根据具体业务调整增强策略、后处理逻辑和类别权重。RGB三通道图像的分割核心工作量从来不在网络结构而是在数据链路。把这层窗户纸捅破后面自然顺畅。