在医学影像分割这个圈子里，3D UNet差不多是绕不开的入门必修课。相比2D UNet，它直接吃三维体数据，能把CT、MRI序列当成一个整体来学，保留了解剖结构在第三维上的连续性。而一份3D UNet的Keras工程，读代码时最容易一头扎进model.py里看网络结构，却忽略了一个事实：真正决定模型能不能收敛、收敛多快的，其实是train.py这个入口脚本。
这篇博客就是要把train.py从头到尾拆开来讲，适合刚接触医学影像分割、或者已经跑通了2D分割想挑战3D的开发者。我会带着你逐段解读训练脚本里的数据流、模型编译、回调配置、超参选择，同时把我自己实战中踩过的坑也一并列出来。读完你不仅能看懂train.py里每一行在干什么，还能照着改出一套属于自己的3D训练流程。
1. 项目整体架构与train.py的定位
1.1 一份典型3D UNet工程的目录结构
一个完整可跑的3D UNet项目，文件划分通常是有套路的。我之前看过并参与过好几个医学影像分割工程，目录风格大同小异：
medseg_project/ ├── data_generator.py # 三维数据读取、预处理、在线增强 ├── model.py # 3D UNet网络结构定义 ├── train.py # 训练入口脚本 ├── predict.py # 推理与后处理 ├── utils/ │ ├── metrics.py # dice、IOU等评估指标 │ └── losses.py # 损失函数定义 ├── config.py # 全局超参数配置 ├── data/ │ ├── train_volume/ │ └── train_label/ └── models/ # 训练产出的权重文件train.py从文件名就能看出，它是整个工程的总调度。它本身不写卷积、不画网络结构，但它会把三个东西串起来：训练数据从哪来、网络结构是谁、用什么指标评判好坏。
从架构上理解，train.py有点像工程队里的现场总包。挖地基的是data_generator，搭结构的是model.py，省得多的是model_checkpoint，但现场指挥、资源分配、处罚规则都是train.py在管。读完train.py，你就掌握了这个工程的核心编排逻辑，其他脚本都是给它服务的。
1.2 为什么训练逻辑要单独抽成一个脚本
很多初学者喜欢脚本套脚本，把读数据、建模型、训练全写在一个大文件里。前两天有位初学者给我看他的代码，一个文件八百多行，数据增强和loss定义混在一起，改一个batch_size要滚动半天才能找到。这种写法在跑通demo时没问题，但进了真实项目就会很痛苦。
把训练逻辑单独抽成train.py，核心不是为了凑文件数量，而是为了隔离变化。数据增强策略、损失函数、网络结构，是三个会独立变化的维度。今天你要把Dice Loss换成Focal Loss，明天的训练数据要从单模态变成多模态，后天的backbone要从UNet换成残差结构。如果这些耦合在一个文件里，每改一个方向都要把另外两块的代码也摸一遍，极易改出隐性bug。
另外还有一个实操原因：train.py通常是团队里多个成员都会碰的文件。算法工程师会调它，做数据清洗的同事也要在交付前把参数改成最优配置。逻辑集中、入口清晰，协作时就很少出现"你改了我那块代码"的冲突。
2. 数据加载与预处理逻辑解析
2.1 三维数据在Keras里的输入形状
先明确一个基础点：在Keras中，3D UNet的输入张量不是三维数组，而是五维张量。以Keras常用的channels_last模式为例，一个输入样本的形状是：
(batch_size, height, width, depth, channels)这里height、width是二维平面尺寸，depth是切片数量。比如一个CT序列在z轴方向有120层，每层是512x512，单通道，那么一个样本就是(512, 512, 120, 1)。真正喂给网络时通常要加batch维度，变成(4, 512, 512, 120, 1)。
这一点特别容易栽跟头。二维图像时代大家的思维定式是(width, height, channels)，三维多了个depth维度，很多人在维度顺序上翻车。Keras的底层卷积层是按照(kernel_dim1, kernel_dim2, kernel_dim3)这样的顺序匹配维度的，如果你的数据维度顺序和网络定义不一致，训练时常常会出现"维度从(2, 128, 128, 128, 1)变成(128, 2, 128, 128, 1)"这种非常隐晦的错误，表面看不报错，指标却一直乱跳。
所以读train.py第一步，先确认数据预处理出的张量形状，再回看model.py的input_shape是否对齐。通常工程会在config.py里维护一个全局的target_shape，train.py负责调用数据脚本时把它传进去，保证两边一致。
2.2 自定义Generator的典型写法
医学影像分割里，label通常是逐体素标注的mask，数据量又大，Keras自带的keras.preprocessing.image.ImageDataGenerator是给二维图片用的，处理不了NIfTI这种三维医学格式。所以工程里几乎都会实现一个继承keras.utils.Sequence的Generator类，然后train.py通过fit_generator / fit接口去消费它。
一个基三本骨架的Generator是这样的：
class SegDataGenerator(keras.utils.Sequence): def __init__(self, sample_list, label_list, batch_size=4, dim=(128, 128, 128), n_channels=1, n_classes=2, shuffle=True): self.sample_list = sample_list self.label_list = label_list self.batch_size = batch_size self.dim = dim self.n_channels = n_channels self.n_classes = n_classes self.shuffle = shuffle self.on_epoch_end() def __len__(self): return int(np.floor(len(self.sample_list) / self.batch_size)) def __getitem__(self, index): indexes = self.indexes[index * self.batch_size:(index + 1) * self.batch_size] batch_x = np.zeros((self.batch_size, *self.dim, self.n_channels), dtype=np.float32) batch_y = np.zeros((self.batch_size, *self.dim, self.n_classes), dtype=np.float32) for i, idx in enumerate(indexes): x, y = self._load_and_preprocess(idx) batch_x[i] = x batch_y[i] = y return batch_x, batch_y def on_epoch_end(self): self.indexes = np.arange(len(self.sample_list)) if self.shuffle: np.random.shuffle(self.indexes)train.py中实际使用它时，会结合config里的参数来实例化：
from data_generator import SegDataGenerator from config import train_cfg train_gen = SegDataGenerator( sample_list=train_samples, label_list=train_labels, batch_size=train_cfg['batch_size'], dim=train_cfg['target_size'], n_classes=train_cfg['num_classes'], shuffle=True )这里有个细节值得展开：为什么用Sequence而不用普通生成器函数？因为Sequence支持多进程安全取数，配合fit_generator的workers参数可以并行加载数据；在显存充足但磁盘IO吃紧的场景下，这个性能差距非常明显。
2.3 预处理里的窗宽窗位与归一化问题
三维医学影像的预处理，和自然图像完全不同。自然图像直接除255归一化就行，但CT这类数据的像素值代表的是组织对X射线的衰减系数，用HU单位表示，范围通常在-1000到+3000左右。直接拿原始值做归一化，网络绝大多数精力都会花在拟合背景和空气上，软组织对比度反而被压缩。
训练脚本里最常做的两步是：窗宽窗位裁剪、再线性归一化。比如处理腹部CT时，会把HU值裁剪到[-100, 240]的区间，这个操作就是把窗位调到40、窗宽拉到340左右，专门突出软组织；然后把这个区间线性映射到[0,1]。
def preprocess_ct(volume): volume = np.clip(volume, -100, 240) volume = (volume - (-100)) / (240 - (-100)) return volume.astype(np.float32)这个窗口参数不是拍脑袋决定的，是放射科常用的显示参数迁移过来的。不同部位有不同默认值，脑部、肺部、骨窗参数都不同。训练前花十分钟确认一下你的数据分布，画个直方图，比调三天学习率都管用。
2.4 在线增强与patch采样
3D体数据往往尺寸巨大，直接全图输入GPU基本撑不住。常规做法是从大体积中裁剪出固定尺寸的子块，也就是patch，再用patch训练。patch大小的选择，要在网络感受野、显存占用、上下文信息三者之间取平衡。常见的做法裁成(128,128,128)或(96,96,96)。
patch从哪里采，是个讲究活。如果只在病灶附近采样，网络训练集分布和推理分布就不一致；如果随机全域采样，很多patch可能完全落在背景区域，dice长期不涨。我之前跑过一个肝脏分割任务，最初随机采样出来的patch大约七成图上根本没有肝，网络学了一堆肝背景的特征，精确率上去了，敏感度却一直起不来。
后来改进方案是混合采样：一部分patch由带标签的mask生成候选框，在候选框附近做随机偏移裁剪；另一部分patch在全图随机采样，用来保持背景分布。两部分的配比一般按7:3，效果立刻好转。train.py里通常不会写采样的具体代码，它调用的是data_generator里实现的采样器，但你要知道，训练效果的瓶颈常常就在这一步。
3. 模型构建与编译细节
3.1 从model.py导入3D UNet结构
train.py里通常情况下就是一行import：
from model import build_3d_unet但真正编译之前，有几个环境层面的细节必须确认。Keras的全局数据格式要和你的输入数据一致。很多老项目跑在Keras 2.x上，这时要检查keras.backend.image_data_format()的值：
import keras.backend as K print(K.image_data_format())如果是channels_first，你的输入张量结构就完全反过来了，shape要变成(n, channels, depth, height, width)。这部分不确认清楚，后面所有维度的代码全白写。
回到模型构建，模型函数一般接受input_shape和num_classes作为参数：
input_shape = (train_cfg['target_size'][0], train_cfg['target_size'][1], train_cfg['target_size'][2], 1) model = build_3d_unet(input_shape=input_shape, num_classes=train_cfg['num_classes'])模型内部就是经典的编码器-解码器结构，编码器逐层下采样提取高层语义特征，解码器逐步上采样恢复空间分辨率，中间用跳跃连接把不同尺度的低级特征拼进来。这个框架在train.py里不需要关心太多，它只是作为一个待训练的黑盒被加载。
3.2 损失函数和评估指标的选择
编译这一步是train.py里信息密度最高的位置。直接看代码：
from utils.losses import soft_dice_loss, weighted_cross_entropy from utils.metrics import dice_coefficient model.compile( optimizer=keras.optimizers.Adam(learning_rate=1e-3), loss=soft_dice_loss, metrics=[dice_coefficient] )很多新手会问，为什么分类任务常用的categorical_crossentropy到分割任务就不香了。问题出在类别不平衡上。以器官分割为例，整张CT里器官体积占比往往只有百分之几，剩下的全是背景。对交叉熵来说，背景像素的梯度贡献会淹没前景，模型很快就学会"全预测背景"走捷径，dice自然上不去。
医学分割领域最常用的替代方案是Soft Dice Loss。Dice系数的计算公式是2倍预测与标签的交集除以两者像素数之和，本质上是直接优化两个分割mask的重叠度。把网络输出概率图参与计算，得到的就是可微的soft dice loss。它天然对类别不平衡不敏感，因为公式的分子分母量纲一致，背景再大也只是把分母撑大而已。
实际工程里更常用的做法是混合loss，把soft dice loss和带权交叉熵按一定权重相加。比如：
def combined_loss(y_true, y_pred): return 0.6 * soft_dice_loss(y_true, y_pred) + 0.4 * weighted_cross_entropy(y_true, y_pred)这样做是让交叉熵在高置信区域提供稳定的梯度，而dice loss在边界区域提供更细致的形状约束。组合权重需要实验调整，但从0.5:0.5起步是很多项目的基础配置。
评估指标的选择也有讲究。训练日志里metric字段用dice_coefficient，而不是keras自带的accuracy。原因很简单，在背景占比极高的分割任务里，accuracy随随便便就能刷到99%，毫无参考价值。Dice才是临床比较关心的体积重叠率，这也是train.py日志里最需要盯住的数字。
3.3 优化器与学习率策略
Adam是训练分割模型最常见的起点。它开了自适应学习率，前期收敛快，对新手友好。靠得住的经验是先把初始学习率设为1e-3或5e-4，跑几十个step观察loss曲线的下降趋势，如果震荡剧烈就降一档到1e-4，如果下降平稳但这个速度太慢，下个阶段再调回1e-3。
学有余力的话，把优化器从Adam换回SGD+momentum也值得一试。SGD在批量数据足够大、训练时间足够长的条件下，最终收敛精度常常比Adam更好，因为它的参数更新路径更平滑。具体代码：
sgd = keras.optimizers.SGD(learning_rate=1e-2, momentum=0.9, nesterov=True)注意SGD的初始学习率一般要比Adam高一到两个量级。这个坑我踩过，直接拿1e-3跑SGD，loss纹丝不动，还以为是代码写错了。
4. 训练循环与回调机制逐行解读
4.1 fit_generator / fit流程中的关键参数
Keras 2.x时代用的是model.fit_generator，Keras 3统一到model.fit。无论哪个版本，关键参数的含义都一样。实际train.py里的调用类似：
history = model.fit( train_gen, epochs=train_cfg['epochs'], validation_data=valid_gen, steps_per_epoch=len(train_gen), validation_steps=len(valid_gen), callbacks=callbacks, verbose=1, initial_epoch=0 )steps_per_epoch要不要显式指定，取决于fit接口内部如何解释数据。如果你传的是Sequence对象，不传它也能通过len(sequence)自动算。但显式写上有个好处：当训练样本数不是batch_size整数倍时，你可以选择扔掉尾部不足一个batch的数据，或者保留它但最后一个小batch的shape和前面不同。为了省事和稳定，工程里常常让steps_per_epoch就取len(train_gen)，宁可丢几条尾巴也不让模型在训练中途报shape错。
epochs设多少，在实际工程中不是拍脑袋定的。真实场景往往配合回调里的EarlyStopping来动态停止，你可以设个大值比如500，让模型和回调自己商量什么时候结束。这样做的好处是不同规模的实验不用频繁改代码。
4.2 ModelCheckpoint、TensorBoard等回调配置
train.py里一个容易写错但极其重要的部分是回调。最核心的ModelCheckpoint：
checkpoint = keras.callbacks.ModelCheckpoint( filepath='./models/best_model.h5', monitor='val_dice_coefficient', mode='max', save_best_only=True, save_weights_only=False, verbose=1 )这里monitor参数决定哪个指标用来判断"最好"。列表里的mode要配套设置，如果监控的是loss，mode就是min；监控dice的话就是max。有个很隐蔽的坑是：如果你自定义的metric名称里包含特殊字符，或者返回的是多输出列表，回调里可能匹配不上，导致save_best_only永远无法生效。排查方法很简单，先把monitor改成'val_loss'跑一轮，确认权重在正常保存，再切回自定义指标。
再看另外三个常用回调的配置：
reduce_lr = keras.callbacks.ReduceLROnPlateau( monitor='val_dice_coefficient', mode='max', factor=0.5, patience=10, min_lr=1e-6, verbose=1 ) early_stop = keras.callbacks.EarlyStopping( monitor='val_dice_coefficient', mode='max', patience=30, restore_best_weights=True ) tensorboard = keras.callbacks.TensorBoard( log_dir='./logs', update_freq='epoch' )ReduceLROnPlateau是学习率策略中比较省心的一种：当监控指标连续patience个epoch不涨时，学习率就降一半。factor取0.5是一个保守又有效的值，太激进容易直接跳出潜在的最优点。EarlyStopping的patience要大于ReduceLROnPlateau的patience，不然模型还没等到学习率降下来的二次收敛窗口，就先被喊停了。restore_best_weights=True这个参数值得强调，它保证训练结束后恢复的是验证集上最优的权重，而不是最后一批epoch的权重，这一点对最终提交模型至关重要。
4.3 训练日志监控与断点续训
训练开始后，tensorboard日志里的loss曲线和val_dice_coefficient曲线是最直观的信息源。我自己习惯在tensorboard里同时盯三样东西：训练loss是否在稳步下降，验证dice是否在同步上升，二者之间的差距有没有急剧扩大。前两样正常但第三样出现，就意味着过拟合来的比预期早，需要赶紧加正则化或者数据增强。
断点续训这个需求在实际训练长模型时几乎是必然出现的。3D医学影像数据量大，一个模型训练三五天很正常，中途机器掉电、杀进程等都是现实风险。续训的方式非常简单：
if train_cfg.get('resume_from'): model = keras.models.load_model( train_cfg['resume_from'], custom_objects={ 'soft_dice_loss': soft_dice_loss, 'dice_coefficient': dice_coefficient } )这里必须强调custom_objects的作用。如果你的模型使用了自定义loss函数或自定义metric，那么加载编译信息时，Keras必须能在全局名称空间里找到这些对象名。不传custom_objects，或者train.py里import的路径和权重文件里记录的路径不一致，加载时会直接抛ValueError。这个问题在团队协作中尤其常见，一个人写loss，另一个人训练，命名对不上就翻车。
续训还有一个注意项：如果配置里改了优化器或学习率，加载的权重只还原网络参数，优化器状态和新学习率是重新初始化的。所以续训前先想清楚，你到底是想严格从上次状态继续，还是只想在已有模型基础上拿新策略重新训后面这一段。
5. 训练过程中的常见问题与排查实录
5.1 显存溢出、数据加载慢等环境类问题
3D UNet最常遭遇的实际问题，就是显存溢出。一个batch设为4、patch尺寸(128,128,128)的模型，显存占用轻松超过11GB。很多开发者倒在这步，以为代码写错了，其实是显存物理容量不够。
排查和解决思路，按损耗程度从小到大排列：
- 调小batch_size，从4降到2甚至1，这一般能把显存峰值砍掉一半。
- 调小patch尺寸，例如从128降到96，显存占用会跟着体素总量呈立方级下降。但要注意感受野会变小，影响大结构的分割效果。
- 检查网络里每个卷积层输出通道数是否可以压缩，比如把第一层的16个通道降到8。
- 检查是否开了cudnn.benchmark等与精度无关但有影响的开关。
- 必要时换显存更大的硬件或者用多卡跑。
数据加载慢是另一个高频问题。判断标准很简单：训练时nvidia-smi里GPU利用率很低，但训练日志里每个epoch的耗时很长。这说明GPU在等数据，瓶颈在磁盘IO或CPU预处理。解决的常规手段是调大fit里的workers和use_multiprocessing：
model.fit( train_gen, workers=8, use_multiprocessing=True, max_queue_size=20 )这里有个现实限制：workers数值不是越大越好。机器CPU核数有限，进程间还有数据拷贝的开销；另外如果你的预处理里用了nibabel等库，多进程模式下每个worker都会载入一份库，内存压力会跟着上来。稳妥做法是先拿workers=4和workers=8各跑一个epoch对比，看耗时有没有显著改善。
5.2 模型不收敛或dice波动大怎么排查
训练日志里的dice系数忽上忽下，是很常见的事故现场。不要急着去改网络结构，先按下面的优先级排查：
第一，确认标签是否正确对齐。三维分割数据的输入和mask来自两个不同的数组，经过裁剪、旋转、翻转后，二者是否依然保持像素级对齐？这个错误非常隐蔽，因为它不报错，只让你训练指标一直在预期值附近徘徊。排查方法是把预处理后的某个batch里，输入数据和mask叠加可视化，眼睛过一遍，比写代码检查更快。
第二，确认one-hot编码的shape对不对。比如二分类任务，Keras的softmax输出层期望的监督标签是两个通道的one-hot编码，如果你直接传了单通道的0/1 mask进来，网络可能在某个不匹配点上悄悄报错，或产生永远无法收敛的梯度。打印batch_y的shape，确认它和网络最后一层输出的通道数一致。
第三，确认损失函数里没有NaN。dice loss的分母是预测值和标签的积之和，这类全局归一化的loss在logits全零或全一时可能出现除零问题。很多工程里会给分母加一个平滑项，比如：
def soft_dice_loss(y_true, y_pred, smooth=1e-6): intersection = K.sum(y_true * y_pred, axis=[1,2,3,4]) union = K.sum(y_true, axis=[1,2,3,4]) + K.sum(y_pred, axis=[1,2,3,4]) dice = (2.0 * intersection + smooth) / (union + smooth) return 1.0 - K.mean(dice)smooth的作用就是防止分母为0，这个值在数值计算里属于最有存在感的1e-6。
第四，确认是否真的在学习。训练前几个epoch里，loss下降一点点其实是正常的，但如果前100个step内部loss纹丝不动，优先把学习率调大一个量级试跑；如果一调大就发疯式震荡，就说明梯度信号本身有问题，回到第一步查数据。
5.3 断点续训和模型产物管理的细节
训练完成后，models目录下可能同时躺着几十个epoch权重文件。业务模型管理上有个实用习惯：把所有训练产物统一命名规则，比如model_epoch{epoch:02d}valdice{val_dice:.4f}.h5，这样从文件名就能看出哪一个对应哪一轮验证集dice。ModelCheckpoint里filepath可以直接用这种方式，Keras会动态替换花括号变量。
文件太多时，还建议定期做一轮清理，只保留最高精度的几个权重，因为磁盘空间被成堆权重文件占满是迟早的事。续训或推理时，优先加载save_best_only对应的best_model.h5。
另外，多卡训练时要注意权重保存的方式。Keras的多GPU训练，model本身被包装成了parallel_model，保存的权重是并行包装层的名字前缀。推理时如果想用单卡加载，可能需要load_weights之后剥掉前缀或者在保存前把权重转成单卡格式，不然会卡在模型结构不匹配的问题上。很多项目简化处理，直接把多卡训练当做并行加速手段，保存的是每个epoch后并行模型里主模型的权重，这个细节在跨机器复现时特别容易出问题。
6. 从train.py延伸出去的三个思考
6.1 训练脚本是复制出来的，但调参过程是不可复制的
train.py写得好，意味着它足够清晰、参数集中、扩展方便。但真正从这份脚本里学到东西的，不是那些代码行，而是在调参过程中积累的判断力。我在实际参与的不同医疗分割项目里，几乎每个项目都会对同一份基础train.py做大量修改，包括采样策略、数据增强、loss构成、学习率调度。真正跑通一个任务的经验沉淀，不是那只改好的py文件，而是你在每个失败实验中识别出来的模式：什么样的错误会导致dice不涨，什么样的数据会让loss震荡。这些东西，恰恰是任何代码注释都写不出来的。
6.2 读train.py时理解数据流比看懂每一行更重要
刚接触源码解析时，很多读者会把注意力放在每一行的语法上，希望一行行解释。但train.py这类编排脚本，真正的阅读方式是跟着数据走一条线：原始NIfTI文件从磁盘进来，经过Generator的裁剪、归一化，变成一个五维张量；这个张量被卷积网络接收，输出一个同样维度但通道数不同的概率图；概率图与真值的mask计算loss，反向传播更新权重；每隔几个epoch回调决定是否存盘、是否降学习率。把这条主线抓住了，train.py里每一行代码在串联这条线的哪个位置，自然而然就清晰了。
6.3 工程化思维才是训练脚本的灵魂
train.py最终交付的，不是一段能跑的代码，而是一套可复现、可扩展、可管理的训练方案。一份好的训练脚本，应该让一个新加入的成员在没有口头交接的情况下，也能看懂超参配置放在哪、数据读取从哪改、网络结构怎么替换、日志怎么看。达到这个标准，训练脚本就成了团队资产。反过来，如果脚本只有自己看得懂，那它充其量只是临时的实验草稿，换个环境和数据就彻底归零了。
我在写自己的train.py时，有一个持续坚持的改法：把config和脚本主体彻底分离。所有会按项目变化的参数，例如patch尺寸、类别数、学习率、训练路径，全部扔进config.py里集中维护。train.py只负责读配置、组装组件、执行训练。这样改完之后，换一个数据集训练时，基本不用改动训练脚本本身，只需要另写一份config或者动态传参。整个工程的可复用性，一下就上来了。
