企业官网建设流程全解析

如果你手头有一个医学图像分割任务，跑通了nnU-Net（v1）的完整流程，拿到了不错的Dice，接下来想往这个"老而稳"的框架里加一点自己的东西——比如把原始图像和血管增强图一起作为输入、把默认U-Net编码器换成更适合落地的轻量结构、或者针对小目标结构重新设计损失函数——那么恭喜你，你进入了"老框架扩展"这个经典的深水区。这篇文章不聊入门教程，专门聊我在扩展nnU-Net v1过程中用到的思路、改过的代码位置，以及那些让模型"看起来在训练，实际上已经废掉"的坑。

先说结论：nnU-Net v1虽然被官方新版本替代，但它的代码结构非常"教科书化"，训练逻辑清晰、接口稳定、第三方程现成，非常适合做二次开发和算法验证。很多人一上来就想直接换框架，没必要。只要你知道扩展的边界在哪里、改哪里、怎么验证，v1反而比很多流水线更顺手。

1. 扩展v1前，先想清楚要扩展什么

1.1 v1留给我们的扩展空间

nnU-Net v1的核心优点不是网络结构本身有多强，而是它把"数据集分析、预处理策略选择、网络结构配置、训练方案、推理和后处理"打包成了一套可复现的流水线。你扩展它，本质上是把流水线里某一个环节换成自己的实现。常见可替换点有三个：数据加载与预处理、网络主体、训练目标与后处理。

我见过不少团队把v1当黑盒用，跑通五折交叉验证就完事了。但一旦要扩展，就必须把黑盒打开。比如v1的Trainer会读取预处理阶段生成的plans.pkl，其中包含patch_size、batch_size、target_spacing、architecture等关键参数。你的扩展如果在网络层面改了输入通道数，却不改plans，训练时DataLoader和网络第一层卷积就对不上。这就是典型的"扩展点没找对"。

所以第一步不是写代码，而是画一张你现在这个任务的流水线图：数据从哪来、预处理做了什么、网络输入是什么形状、损失怎么算、预测完怎么还原到原始分辨率。把这张图挂在屏幕上再动手。

1.2 "扩展"不等于"重写"

扩展v1的一个常见误区是把整个框架的代码都改成自己的风格。我早期也干过这种事：为了加一个trick，把nnUNetTrainerV2整个复制出来，改了十几处，最后训练倒是能跑了，但推理阶段忘记同步修改，结果一评估全崩。

正确的做法是"最小侵入"。v1本身设计了继承机制，你完全可以创建一个新的Trainer类继承官方类，只重写需要变化的方法。比如要换损失，就重写initialize_loss_function；要换网络，就重写initialize_network；要保持训练策略不变，其他方法全部继承。这样扩展点集中、可对比、也容易回滚。

我这里有一个经验原则：凡是官方已经替你验证过的逻辑，不要动它；凡是和你研究目标相关的逻辑，尽量通过继承重写。

1.3 什么情况下别用v1扩展，直接换v2

v1适合扩展的场景是：你已经跑通了现有模型，想在训练策略、损失、输入模态、后处理这些上层模块上做实验。v1的代码简单，3000多行你能全部读完，出问题好定位。但如果你的目标是彻底替换主干网络为Transformer结构、大规模分布式训练、或者是跨数据集联合训练，那v1的底层设计会成为阻碍，老老实实去用v2更快。

我不是版本党，但从工程效率上讲，"能扩展就扩展，不能扩展就换"才是正确心态。

2. 扩展前必须吃透v1的训练骨架

2.1 三个环境变量和一个Task目录

扩展nnU-Net v1，第一步不是改代码，而是理解它如何管理数据。旧版v1依赖三个环境变量：nnUNet_raw_data_base、nnUNet_preprocessed和nnUNet_results。官方文档里给出的初始化大概是这样的。

export nnUNet_raw_data_base="/path/to/nnUNet_raw_data" export nnUNet_preprocessed="/path/to/nnUNet_preprocessed" export nnUNet_results="/path/to/nnUNet_trained_models"

原始数据目录下面必须按照TaskXXX_Name的规则组织，例如模拟项目"某肝脏血管分割任务"的目录结构如下。

Task501_LiverVessel/ ├── imagesTr/ │ ├── case001_0000.nii.gz │ ├── case002_0000.nii.gz ├── labelsTr/ │ ├── case001.nii.gz │ ├── case002.nii.gz ├── imagesTs/ │ ├── case101_0000.nii.gz └── dataset.json

imagesTr里的文件名后缀_0000表示第一个输入模态。如果你要扩展成多通道输入，比如"原始CT + 血管增强图"，那文件名就是case001_0000.nii.gz和case001_0001.nii.gz。v1的DataLoader会按模态编号自动把多个nii文件堆叠成多通道输入，这个机制对扩展非常友好。

2.2 plan_and_preprocess 在替我们做哪些决定

v1的前置命令nnUNet_plan_and_preprocess看起来是个黑盒，但它生成的文件决定了你后面所有扩展的口径。它的作用是分析整个训练集：统计体素间距、强度分布、类别比例，然后自动决定target_spacing、patch_size、batch_size和网络拓扑。这些信息最终存在preprocessed目录下某个Task子目录里的plans.pkl和dataset_properties.pkl中。

扩展多通道输入时，很多人只改了文件命名，却忘了重新执行预处理。旧版v1的预处理阶段默认会对每个模态分别计算z-score归一化，新增的模态如果强度分布完全不同（比如增强图是梯度幅值，分布高度偏斜），按默认方式归一化可能不合理。你需要检查nnUNetTrainerV2的预处理设置，针对新模态单独调整归一化策略。

另一个隐蔽点：v1在预处理时会做重采样到target_spacing。如果你的扩展目标包含"在原始高分辨率下训练"，那你要么修改规划阶段的策略，要么在推理阶段做resample时对齐。我见过不少扩展项目在训练时用了新分辨率，推理时却忘了把预测结果重采样回原始间距，结果HD95惨不忍睹。

2.3 训练、推理、后处理是三条必须对齐的链路

v1的扩展不能只盯着训练。官方训练命令nnUNet_train和推理命令nnUNet_predict，以及后处理脚本nnUNet_determine_postprocessing，它们之间存在隐含的一致性约束。

比如推理阶段的nnUNet_predict会默认使用训练时保存的plans.pkl来决定如何重采样、如何滑动窗口。如果你在扩展时改了预处理策略，但没有把新策略同步到推理流程，模型等于在一个分布上训练，在另一个分布上预测，性能崩盘是必然的。这个我在第4节详细展开。

所以扩展v1的正确工作流是：先统一口径（计划和预处理），再改训练，然后验证推理，最后才谈后处理。任何一步单独跑通都不算完成。

3. 三个可落地的扩展方向

3.1 扩展方向一：多通道输入

医学图像分割最常见的扩展需求之一，是把多种模态或派生特征同时输入网络。nnU-Net v1本身支持多模态，你只要按_0000、_0001的规则准备数据即可。我实际做过的一个场景是：原始CT图像加上一个血管增强滤波图，两者一起训练，目标是小血管分割。

数据准备阶段，我用简单滤波算子生成增强图，然后和原图一起保存。

import nibabel as nib import numpy as np ct = nib.load("case001_0000.nii.gz").get_fdata() enhanced = vesselness_enhancement(ct) # 自定义血管增强滤波 combined = np.stack([ct, enhanced], axis=-1) # 形状: H, W, D, 2 affine = nib.load("case001_0000.nii.gz").affine nib.save(nib.Nifti1Image(combined.astype(np.float32), affine), "case001_0001.nii.gz")

注意这里要把增强图单独存成_0001文件，而不是和原图拼在一个nii里。原因在于v1的DataLoader和预处理流程默认一个文件就是一个模态，它会对每个模态单独做归一化和数据增强。你拼在一起，等于关闭了框架自带的模态处理逻辑，后续会踩很多坑。

训练阶段，网络第一层卷积的输入通道数会由plans自动设为2。但有几个雷区：第一个雷区是新增模态的数值范围，v1默认z-score归一化，血管增强图的灰度分布往往极度右偏，log变换会更合理，这需要你重写预处理组件；第二个雷区是数据增强时几何变换要对所有模态同步，v1的多模态加载机制天然保证这一点，但如果你手动拼通道，增强就会把两个模态当成一个整体来插值，信息会混叠。

3.2 扩展方向二：换一个更轻的编码器

第二个常见扩展方向是网络结构改造。v1默认的3D U-Net在显存和推理时间上对不少落地场景不太友好。我试过把编码器的特征通道数整体缩小，默认配置第一层是32个通道，我改成16，后面每层翻倍模式不变。

实现方式依然建议继承官方Trainer。

from nnunet.training.network_training.nnUNetTrainerV2 import nnUNetTrainerV2 from nnunet.network_architecture.generic_UNet import Generic_UNet class MyLightUNetTrainer(nnUNetTrainerV2): def initialize_network(self): # 读取官方plan里的基础结构配置 self.network = Generic_UNet( input_channels=self.num_input_channels, base_num_features=16, # 默认是32，减半 num_classes=self.num_classes, num_pool_per_axis=self.pool_op_kernel_sizes, conv_op=nn.Conv3d, # 2D场景改为nn.Conv2d norm_op=nn.InstanceNorm3d, ... )

这里最需要注意的是Generic_UNet的构造参数必须和plans.pkl中规划的网络结构对齐。你不能只改base_num_features，还要关注深度的变化带来感受野变化。如果层数减少，小目标可能分割得更细，但大目标的上下文信息不够。我在某个模拟项目上把深度从6层降到5层，整体Dice掉了接近2个百分点，不甘心又去调了patch size，才算追回来。

另一个坑是预训练权重。如果你要用ImageNet或自监督预训练权重来初始化轻量编码器，v1的加载代码是严格匹配state_dict的，一旦通道数变化，直接加载会报错。这种情况要么自己写一个映射函数，只加载匹配的层，要么在initialize_network里直接跳过预训练，从头训。我的建议是：分割任务数据量够的话，先不要花时间调预训练权重，v1的初始化策略已经够用了。

3.3 扩展方向三：自定义损失函数与后处理

v1官方默认用的是DiceLoss和CrossEntropyLoss的组合。这个组合对大多数器官分割都很稳，但如果你的目标是小物体、极度不平衡的类别，或者边界模糊的结构，默认损失就不一定够用。

扩展损失函数，我推荐重写initialize_loss_function。

import torch.nn as nn from nnunet.training.loss_functions.dice_loss import DC_and_CE_loss class MyLossTrainer(nnUNetTrainerV2): def initialize_loss_function(self): self.loss = DC_and_CE_loss( {'batch_dice': True, 'smooth': 1e-5, 'do_bg': False}, {'weight': torch.Tensor([0.1, 1.0]).cuda()} )

这种继承方式的好处是，你还能保留原来的训练循环、学习率调度和验证逻辑。如果换成Focal Loss，就要注意数值稳定性，v1默认开了混合精度（AMP），Focal Loss里的torch.log在低精度下容易变成NaN，我第4节会讲具体表现和解决办法。

后处理方面，v1默认的nnUNet_determine_postprocessing会基于训练集五折交叉验证结果决定是否保留最大连通域、是否移除小区域。扩展后处理最常见的需求是"剪切掉某些预测区域"或"按临床规则合并标签"。我的经验是，不要自己去改官方后处理脚本，而是训练完之后，把softmax概率图保存成nifti文件，再写独立的脚本做规则筛选。这样每一步都可视化、可验证、可回滚。

4. 扩展踩坑实录：看起来在训练，实际已经废掉

4.1 预处理不一致：训练分布变，部署全崩

这是我反复踩、也看别人反复踩的一个大坑。场景是这样的：我给某个分割任务加了新的模态，重新预处理后训练Dice涨了不少，但一跑推理，结果比旧版差得多。查了整整两天，才发现问题出在归一化。

原来v1的预处理阶段会把每个模态的“规范化参数”记录在plans.pkl里。训练和推理都要依赖这些参数。而我为了图省事，直接在原图nii上叠加了一个自定义预处理，绕过了v1的数据加载器，结果训练时网络看到的是归一化后的新分布，推理时我喂进去的又是另一种分布。

正确的做法是：扩展预处理时，要么统一在DataLoader里完成所有变换，要么在plan_and_preprocess阶段就把新模态落地成预处理后的nii文件。绝不能在训练和推理之间搞两套预处理逻辑。我自己后来规定：任何扩展，只要涉及输入数据形态，必须走一遍完整的plan_and_preprocess流程，重新生成预处理文件，再训练。

4.2 网络替换导致的状态字典错位

场景是：我想用轻量编码器加载官方预训练权重，结果nnUNet_load_pretrained_weights报错说键不匹配。原因很简单——我把base_num_features从32改成了16，编码器第一层卷积的权重形状从32 x input_channels x ...变成了16 x input_channels x ...，预训练模型的键对不上。

如果你真的想加载预训练权重，一个可行的做法是手动加载，忽略不匹配的层。

pretrained = torch.load("model_final_checkpoint.model.pkl")["state_dict"] model_dict = model.state_dict() for k, v in model_dict.items(): if k in pretrained and pretrained[k].shape == v.shape: model_dict[k] = pretrained[k] else: print(f"skip {k}") model.load_state_dict(model_dict)

不过我想提醒的是，v1默认训练流程本身没有预训练，也从零开始训出了不错的结果。除非你的任务数据量很少、且你的网络结构和官方非常接近，否则花时间去调预训练权重性价比很低。我见过太多人卡在"权重加载成功但效果更差"的怪圈里，不如直接让数据说话。

4.3 AMP精度与自定义Loss的兼容性

v1的Trainer默认使用AMP（自动混合精度），实现细节里有一个grad_scaler。如果你新写的Loss包含torch.log(1 + x)、torch.exp(-x)这类指数对数操作，在半精度下可能因为梯度溢出输出NaN。训练过程看起来loss一开始就很高，随后掉到0并卡住，但验证集Dice始终是0。

我采取的办法是，在自定义Loss的forward内部，把输入显式转成float32再计算。

class MyFocalLoss(nn.Module): def forward(self, pred, target): pred = pred.float() # 关键：切回float32 ...

同时在Trainer里重写一下混合精度相关的初始化，如果实验不需要AMP，直接关掉也可以。但我要提示：关掉AMP会让训练速度下降，大约慢30%到50%，而且显存占用会上升。所以尽量保留AMP，代价只是在Loss内部做一次float()，影响微乎其微。

4.4 patch_size与显存的隐性失衡

换轻量编码器后，很多人觉得显存变富余了，顺手就把patch_size调大了。这个操作本身没错，但你必须同步改batch_size和num_pool_per_axis。否则网络深度不够、patch size增大后感受野覆盖比重变化，分割精度可能出现局部下降。

比如原来4次下采样覆盖了整个patch，你换了3次下采样的轻量编码器，patch size再翻一倍，最后的feature map上一个像素对应的原图范围就变了。更隐蔽的是，plans.pkl中的patch_size和batch_size是自动规划的，你手动改训练参数时，很多地方从plans里读数，可能造成训练配置不一致。我建议：改任何结构相关参数，都重新打印一遍关键日志，确认训练时实际使用的patch_size、batch_size、num_pool_per_axis和你预期的完全一致。

5. 扩展成果如何落地：评估、集成与止损判断

5.1 用官方评估工具链考核扩展结果

扩展完v1不是训练完就结束了。v1提供了一套完整的评估工具，nnUNet_find_best_configuration和nnUNet_evaluate_folder，你应该把这些流程接到自己的扩展实验里。

find_best_configuration会遍历你训好的各折模型，按验证集Dice找出最优配置，并推荐集成方案。自定义Trainer的结果存放位置要和官方保持一致，否则工具找不到模型。我通常在nnUNet_results目录下新建一个以自己项目命名的子目录，然后把训练好的模型按官方命名规则放进去，这样所有工具链都不用改。

评估指标方面，不要只看Dice。对扩展场景，HD95（95%豪斯多夫距离）和NSD（表面距离）更重要。v1仓库自带的评估脚本已经计算这些指标，直接用即可。

5.2 什么时候该停止扩展v1

这个建议可能有点反直觉：如果你发现一个扩展需要同时修改预处理、网络结构、训练策略和推理脚本四个环节，那说明你已经不在"扩展"v1，而是在“重写”v1。这时候直接切到v2通常更划算。

v2在很多方面做了重构：数据集接口更干净、训练与推理的配置更统一、对自定义网络的支持更友好。但v2的代码复杂度也更高，定制成本不一定低。我的判断标准是：扩展总量控制在官方总代码量20%以内，继续用v1；超过这个比例，迁移v2。这个比例不是硬标准，但每次都能帮我做决断。

另外，如果团队里其他成员没有v1的维护经验，你留下的扩展代码就必须有清晰的文档。我习惯在每个自定义文件头部写清楚"改了什么、为什么改、动了哪个上游函数、如何验证"。扩展老框架最可怕的不是改不成功，而是改成功了但过三个月没人知道哪里改过，一升级就崩。

最后分享一个我个人的习惯：每次给v1做扩展，我都先建立一个"基线分支"，跑通一次官方默认流程并记录全部指标，然后再开新的分支做扩展。新的训练结果低于基线的时候，我会立刻回去检查是不是哪一步改坏了，而不是继续加新模块试图“凑”回来。基线模型是你最忠实的对照组，保存好它，扩展之路才不会变成翻车之路。