自定义模型训练:nnU-Net 插件用户手册

Train Custom Model: nnU-Net - User Manual

Dragonfly Prototype Apps · Train Custom Model: nnU-Net...

版本 Version 1.0 · 2026-07-04

第一部分 中文手册

目录

1. 简介

2. 适用场景

3. 安装与启用

4. 运行环境与首次配置

5. 界面说明

5.1 Setup(环境搭建)分页

5.2 Train(训练)分页

5.3 Infer(推理)分页

6. 使用步骤

7. 参数说明

8. 输出结果

9. 常见问题与故障排除

10. 注意事项与已知限制

11. 参考资料

1. 简介

自定义模型训练:nnU-Net 是一个 Dragonfly 插件,让您直接使用场景中的自有数据训练一个自定义的 nnU-Net v2 三维深度学习分割模型。您为插件提供一个或多个“训练样本”,每个样本由一个图像 Channel(体数据)和一个包含标注类别的 MultiROI 组成;插件会在您的 NVIDIA GPU 上训练模型。训练完成后,可将模型应用(推理)到任意 Channel,预测得到的分割结果会以一个新的 MultiROI 形式发布回场景中,可直接用于测量与可视化。

底层引擎与算法。 本插件封装了开源框架 nnU-Net v2(Python 包名 nnunetv2)——一个“自配置”的医学与体数据分割框架。它会根据您数据的分辨率、体素间距和标注情况自动确定网络结构与预处理流程,因此您无需具备深度学习专业知识。计算后端为 PyTorch(CUDA 版),并使用 SimpleITK 在磁盘上读写 NIfTI(.nii.gz)数据集。

运行隔离。 所有重型计算(PyTorch CUDA、nnU-Net v2、SimpleITK)都运行在插件专属的 Python 虚拟环境(venv)中,通过子进程 + JSON 文件通信驱动,绝不在 Dragonfly 自带的 Python 中导入这些库。这样既避免污染 Dragonfly 环境,也让长时间训练在独立进程中稳定运行。

许可证要点:nnU-Net v2 为开源软件(Apache-2.0 许可),PyTorch 采用 BSD 风格许可,SimpleITK 采用 Apache-2.0 许可。这些第三方组件均在您首次点击 Setup Environment 时从各自官方渠道下载安装到插件专属环境中,并遵循其各自的许可条款。

2. 适用场景

只要您需要针对自有样品定制自动三维分割,本插件都适用。典型场景包括:

• 材料科学与工业 CT: 分割孔隙、裂纹、纤维或不同物相。
• 生命科学: 分割显微 CT、显微镜体数据中的器官、细胞或组织结构。
• 铸件与增材制造(AM)零件: 缺陷检测与分割。

本插件特别适合反复分割同类数据集的用户:只需手工标注少量示例样本作为训练数据,训练出模型后即可对其余大量数据自动分割,大幅减少人工标注工作量。由于 nnU-Net 会自动配置,您无需调整网络结构或编写深度学习代码。

3. 安装与启用

本插件随 Prototype Labs & Apps 完整安装包(Full Package) 分发。安装步骤如下:

1. 解压完整安装包到一个较短的路径(例如 C:\PL\),避免路径过长(Windows 260 字符限制)。

2. 双击运行 `Install_FullPackage.bat`。

3. 在弹出的对话框中选择核心安装模式(Fresh 全新 / Compatible 兼容),并在应用列表中勾选 “Train Custom Model: nnU-Net...”。注意:所有插件默认未勾选,必须手动勾选本插件才会安装。

4. 点击 Install 并等待控制台完成。

5. 完全退出并重启 Dragonfly(菜单只在启动时被扫描)。

重启后,插件出现在菜单 Prototype Apps ▸ Train Custom Model: nnU-Net... 下(位于 “Train Custom Model” 分组,与 LocateAnything、Cellpose 等自定义模型训练器并列)。点击后会打开一个可浮动的面板窗口。

以后修改勾选。 最简单的方式是在 Dragonfly 内操作:打开 Developer ▸ Prototype Labs... ▸ Menu Item Manager,底部 “Prototype Apps (Full Package)” 列表中每个应用一个勾选框——勾选=部署,取消=移除菜单项;重启 Dragonfly 生效。停用绝不会删除已搭建的插件环境(venv),重新启用可立即使用。

重要:安装时不会下载任何深度学习组件。插件的运行环境需在首次使用时于面板内手动搭建(见第 4 章)。

4. 运行环境与首次配置

首次使用本插件前,必须先在 Setup 分页搭建插件专属的运行环境。点击 Setup Environment(build venv + install torch + nnunetv2) 按钮,插件会自动完成以下工作:

1. 使用一个“基础 Python”创建一个专属虚拟环境(venv)。默认使用 Dragonfly 自带的 Python(其标准库 venv 与 pip 均可用),因此您无需另行安装 Python。

2. 从指定的 CUDA 轮子索引安装 torch + torchvision(默认索引 https://download.pytorch.org/whl/cu124)。

3. 安装其余依赖:nnunetv2、SimpleITK、numpy。

4. 运行一次冒烟测试,导入 torch/nnunetv2/SimpleITK 并报告是否检测到可用 CUDA GPU。

下载体积与耗时。 首次搭建大约需要下载 3–6 GB(主要是 CUDA 版 PyTorch),耗时数分钟到十几分钟,取决于网络速度。此过程需要联网。搭建完成后,面板会把 venv 的 python 路径填入 “nnU-Net venv python” 字段并持久保存,下次无需重复搭建。

硬件要求。 训练与推理都需要一块 NVIDIA GPU(支持 CUDA)。3D 全分辨率训练对显存与算力要求较高,训练可能持续数小时。

Turing 架构显卡(如 RTX 8000,sm_75):cu124 轮子可运行,但只能使用 fp16,绝不能使用 bf16(Turing 没有 bf16 张量核)。

环境安装位置。 venv 建在插件安装目录内:%LOCALAPPDATA%\comet\<Dragonfly版本>\pythonUserExtensions\GenericMenuItems\NNUNet\venv。训练产物(数据集与模型)默认写入 C:\nnUNetWorkspace(可在 Train 分页更改)。设置项会持久保存到插件目录下的配置文件,并有一份备用副本位于 %LOCALAPPDATA%\NNUNetCustom\config.json。

失败时的替代方案。 若默认基础 Python 无法创建带 pip 的 venv(例如缺少 ensurepip),可在 Base Python (build) 字段填入另一个可用的 CPython 3.9+ 解释器路径(例如 C:\Python312\python.exe)或形如 py -3.12 的启动器命令,然后重新点击 Setup Environment。若某次搭建中途失败留下了无 pip 的半成品 venv,插件会在下次搭建时自动重建它。若 torch 安装报错,请核对 CUDA 轮子索引是否与您的显卡驱动匹配(例如 cu121 / cu124)。

5. 界面说明

面板顶部有一段蓝色说明文字,下方为三个分页:Setup / Train / Infer;窗口底部有一个绿色结果标签和一个只读的日志文本框,用于显示进度与错误信息。

5.1 Setup(环境搭建)分页

5.2 Train(训练)分页

分为三块:工作区文件夹、① 添加训练样本、② 训练。

5.3 Infer(推理)分页

6. 使用步骤

工作流 A:训练自定义模型

1. 在场景中准备好训练数据:每个样本需要一个图像 Channel 和一个已手工标注好类别的 MultiROI(MultiROI 中的前景类别即为要学习的分割目标)。

2. 打开插件,进入 Setup 分页,点击 Setup Environment 搭建环境(仅首次需要)。

3. 进入 Train 分页,在 Output folder 指定工作区文件夹(默认 C:\nnUNetWorkspace)。

4. 点击 Refresh channels / MultiROIs 枚举场景对象。

5. 从 Image Channel 与 MultiROI (labels) 下拉框中选择一对,点击 Add case;如有多个样本,重复此步逐一添加(Training cases 计数随之增加)。

6. 设置 Configuration(默认 3d_fullres)、Fold(默认 0)与其他选项。

7. 点击 Train。训练进度、当前 epoch、train_loss 与 dice 会显示在日志框中。训练可能持续数小时。

8. 完成后,结果标签显示 “Trained.”,并把结果文件夹路径自动预填到 Infer 分页的 Results folder。

训练开始时,插件会跳过“前景标签数为 0”的样本;若所有样本都没有前景标签,训练会报错。请确保 MultiROI 中确实绘制了标注。

工作流 B:应用模型进行预测

1. 进入 Infer 分页(若刚完成训练,Results folder 已自动填好;否则用 … 按钮选择训练产出的 nnUNet_results 文件夹)。

2. 点击 Refresh 并从 Image Channel 选择要预测的 Channel。

3. 确认 Configuration 与 Fold 与训练时一致。

4. 可选:在 Result title 填写输出 MultiROI 的名称。

5. 点击 Infer -> publish MultiROI。预测完成后,新的 MultiROI 会发布到源 Channel 的网格上,可在对象树中查看。

7. 参数说明

关于 Epochs:nnU-Net v2 的训练轮数固定在其训练器类中(默认 1000),没有命令行开关可修改。因此面板中的 Epochs 字段目前仅用于缩放进度条的显示比例,并不会改变真实训练时长。将其设为 0 即使用 nnU-Net 默认设置。

8. 输出结果

训练输出。 训练在工作区(Output folder)下生成 nnU-Net 标准目录结构:nnUNet_raw(原始数据集,含 Dataset001_Custom,其中 imagesTr 存图像、labelsTr 存标签,均为 .nii.gz 格式,并有一个 dataset.json 描述通道与类别)、nnUNet_preprocessed(预处理产物)与 nnUNet_results(训练好的模型)。推理时需要指向的正是 nnUNet_results 文件夹。

推理输出。 推理成功后,预测得到的分割标签会作为一个新的 MultiROI 发布回场景,叠加在源图像 Channel 的网格上(几何与源 Channel 对齐)。该 MultiROI 名称取自 Result title 字段(默认 nnUNet_prediction)。

如何查看。 在 Dragonfly 对象树(Object Manager)中找到新生成的 MultiROI,勾选可见性即可在 2D/3D 视图中叠加显示。之后可像任何普通 MultiROI 一样对其进行测量、着色和进一步处理。

9. 常见问题与故障排除

问:点击 Setup Environment 失败,提示 venv 创建失败或没有 pip 怎么办?

答:说明所选基础 Python 缺少可用的 venv/pip。请在 Base Python (build) 字段填入另一个 CPython 3.9+ 的路径(如 C:\Python312\python.exe)或 py -3.12 之类的启动命令,再重试。半成品 venv 会被自动重建。

问:训练/推理报错提示 CUDA 不可用或没有 GPU?

答:本插件必须有 NVIDIA GPU。请确认已安装合适的 NVIDIA 驱动,且 Setup 时安装的 torch CUDA 版本与驱动匹配(核对 Torch CUDA wheel index,如 cu121/cu124)。搭建结束的冒烟测试会报告是否检测到可用 CUDA。

问:训练报“显存不足 / out of memory”?

答:3D 全分辨率训练显存占用大。可改用更省显存的 Configuration(如 3d_lowres 或 2d),或减小训练体数据的尺寸,再重试。

问:点击 Train 提示没有训练样本,或提示所有样本前景标签为 0?

答:请先点击 Refresh,选择一对 Channel + MultiROI 并点击 Add case(计数应大于 0)。若提示前景标签为 0,说明 MultiROI 中没有绘制任何标注类别,请先在该 MultiROI 上画出标注。

问:数据集完整性校验(verify dataset integrity)失败?

答:样本数量过少或几何/间距异常时校验可能拒绝数据集。可尝试取消 Verify dataset integrity 复选框,或增加训练样本数量、检查体素间距是否合理。

问:菜单里找不到该插件?

答:请确认安装时已勾选本插件(默认未勾选),并且安装后完全重启了 Dragonfly(菜单仅在启动时被扫描)。仍未出现时,可在 Developer ▸ Prototype Labs... ▸ Menu Item Manager 中勾选后再次重启。

10. 注意事项与已知限制

• 必须联网 + NVIDIA GPU: 首次 Setup 需联网下载数 GB;训练/推理需 CUDA GPU。
• Epochs 字段不改变训练时长: nnU-Net v2 无修改轮数的命令行开关,该字段仅缩放进度条(见第 7 章)。
• Turing 显卡仅用 fp16: RTX 8000 等 sm_75 显卡请勿使用 bf16。
• 训练耗时长: 3D 全分辨率训练可能持续数小时,默认不设硬性超时。
• Configuration/Fold 须匹配: 推理时选择的配置与折号必须与训练时相同,否则会找不到对应模型。
• 单通道图像: 数据集按单个图像通道(channel_names {"0": "image"})构建,背景类别固定为 0。
• Windows 环境: 当前发行版运行模式仅支持 windows;插件内部已禁用 torch.compile 并限制数据增强工作进程数以提升在 Windows 上的稳定性。
• 首次搭建的组件版本由官方最新源决定: nnU-Net v2 的命令行参数对版本较敏感;若上游版本更新导致行为变化,请留意日志中的报错信息。

11. 参考资料

• nnU-Net(nnunetv2)官方仓库与文档:https://github.com/MIC-DKFZ/nnUNet
• PyTorch 官方网站(CUDA 轮子索引):https://pytorch.org/ 、 https://download.pytorch.org/whl/
• SimpleITK 官方网站:https://simpleitk.org/
• 完整安装包的安装/启用/卸载说明:随包 README(Full Package)

Part II English Manual

Contents

1. Overview

2. Use Cases

3. Installation and Enabling

4. Runtime Environment and First-Run Setup

5. Interface Guide

5.1 Setup Tab

5.2 Train Tab

5.3 Infer Tab

6. Step-by-Step Usage

7. Parameter Reference

8. Outputs

9. FAQ and Troubleshooting

10. Notes and Known Limitations

11. References

1. Overview

Train Custom Model: nnU-Net is a Dragonfly plugin that lets you train a custom nnU-Net v2 3D deep-learning segmentation model directly from your own in-scene data. You provide one or more training cases, each consisting of an image Channel (a volume) plus a MultiROI holding the labeled classes; the plugin trains the model on your NVIDIA GPU. Once trained, you can apply (infer) the model to any Channel, and the predicted segmentation is published back into the scene as a new MultiROI, ready for measurement and visualization.

Underlying engine and algorithm. The plugin wraps the open-source nnU-Net v2 framework (Python package nnunetv2) — a self-configuring framework for medical and volumetric segmentation. It automatically derives the network architecture and preprocessing from your data's resolution, voxel spacing and labels, so no deep-learning expertise is required. The compute backend is PyTorch (CUDA build), with SimpleITK used to read/write the NIfTI (.nii.gz) dataset on disk.

Isolation. All heavy computation (PyTorch CUDA, nnU-Net v2, SimpleITK) runs in the plugin's own dedicated Python virtual environment (venv), driven via subprocess + JSON file-based IPC. These libraries are never imported into Dragonfly's embedded Python, keeping Dragonfly's environment clean and letting long training runs proceed stably in a separate process.

Licensing highlights: nnU-Net v2 is open source (Apache-2.0), PyTorch uses a BSD-style license, and SimpleITK is Apache-2.0. These third-party components are downloaded from their official channels into the plugin's dedicated environment the first time you click Setup Environment, under their respective license terms.

2. Use Cases

Use this plugin whenever you need automated 3D segmentation tailored to your own samples. Typical scenarios include:

• Materials science and industrial CT: segment pores, cracks, fibers, or phases.
• Life sciences: segment organs, cells, or tissue structures in micro-CT and microscopy volumes.
• Castings and additive-manufactured (AM) parts: defect detection and segmentation.

It is especially useful for users who repeatedly segment similar datasets: label a few example cases once, train a model, then segment the rest of your data automatically, greatly reducing manual annotation effort. Because nnU-Net configures itself, you need not tune the network or write any deep-learning code.

3. Installation and Enabling

This plugin ships with the Prototype Labs & Apps Full Package. To install:

1. Unzip the Full Package to a short path (e.g. C:\PL\) to avoid Windows' 260-character path limit.

2. Double-click `Install_FullPackage.bat`.

3. In the dialog, choose the core install mode (Fresh or Compatible) and tick "Train Custom Model: nnU-Net..." in the app list. Note: all plugins are unticked by default — you must tick this plugin for it to be installed.

4. Click Install and wait for the console to finish.

5. Fully quit and restart Dragonfly (menus are scanned only at startup).

After the restart, the plugin appears under Prototype Apps ▸ Train Custom Model: nnU-Net... (in the "Train Custom Model" group, alongside the LocateAnything and Cellpose custom-model trainers). Clicking it opens a floating panel window.

Changing your choices later. The easiest way is inside Dragonfly: open Developer ▸ Prototype Labs... ▸ Menu Item Manager; the "Prototype Apps (Full Package)" list at the bottom has one checkbox per app — tick = deploy, untick = remove the menu entry; restart Dragonfly to apply. Disabling never deletes a plugin's built environment (venv), so re-enabling is instant.

Important: no deep-learning components are downloaded at install time. The plugin's runtime environment must be built on first use, from within the panel (see Chapter 4).

4. Runtime Environment and First-Run Setup

Before first use you must build the plugin's dedicated environment on the Setup tab. Click the Setup Environment (build venv + install torch + nnunetv2) button, and the plugin automatically:

1. Creates a dedicated virtual environment (venv) using a "base Python". By default it uses Dragonfly's own bundled Python (whose stdlib venv + pip work), so you need no separate Python installation.

2. Installs torch + torchvision from the specified CUDA wheel index (default https://download.pytorch.org/whl/cu124).

3. Installs the remaining dependencies: nnunetv2, SimpleITK, numpy.

4. Runs a smoke test that imports torch/nnunetv2/SimpleITK and reports whether a usable CUDA GPU was detected.

Download size and time. The first build downloads roughly 3–6 GB (mostly the CUDA build of PyTorch) and takes several minutes to over ten minutes, depending on your connection. It requires internet access. When finished, the panel fills the venv's python path into the "nnU-Net venv python" field and saves it, so you need not rebuild next time.

Hardware requirements. Both training and inference need an NVIDIA GPU (CUDA-capable). 3D full-resolution training is demanding on VRAM and compute, and training can run for hours.

Turing GPUs (e.g. RTX 8000, sm_75): the cu124 wheels run, but use fp16 only — never bf16 (Turing has no bf16 tensor cores).

Where the environment lives. The venv is built inside the installed plugin directory: %LOCALAPPDATA%\comet\<Dragonfly version>\pythonUserExtensions\GenericMenuItems\NNUNet\venv. Training artifacts (dataset and model) default to C:\nnUNetWorkspace (changeable on the Train tab). Settings persist to a config file in the plugin directory, with a fallback copy at %LOCALAPPDATA%\NNUNetCustom\config.json.

Fallback if setup fails. If the default base Python cannot create a venv with pip (e.g. missing ensurepip), enter another usable CPython 3.9+ in the Base Python (build) field (e.g. C:\Python312\python.exe) or a launcher command like py -3.12, then click Setup Environment again. A half-built venv left by a failed run (has python but no pip) is automatically rebuilt on the next attempt. If the torch install errors, check that the CUDA wheel index matches your driver (e.g. cu121 / cu124).

5. Interface Guide

The panel shows a blue note at the top and three tabs: Setup / Train / Infer; at the bottom are a green result label and a read-only log box that display progress and error messages.

5.1 Setup Tab

5.2 Train Tab

Organized into three groups: workspace folder, (1) add training cases, and (2) train.

5.3 Infer Tab

6. Step-by-Step Usage

Workflow A: Train a Custom Model

1. Prepare training data in the scene: each case needs an image Channel and a MultiROI with manually painted classes (the MultiROI's foreground classes are the segmentation targets to learn).

2. Open the plugin, go to the Setup tab, and click Setup Environment to build the environment (first time only).

3. Go to the Train tab and set the Output folder (default C:\nnUNetWorkspace).

4. Click Refresh channels / MultiROIs to enumerate scene objects.

5. Select a pair in Image Channel and MultiROI (labels), then click Add case; repeat for each case (the Training cases count increases accordingly).

6. Set Configuration (default 3d_fullres), Fold (default 0), and other options.

7. Click Train. Progress, current epoch, train_loss and dice appear in the log box. Training may run for hours.

8. When done, the result label reads "Trained." and the results-folder path is pre-filled into the Infer tab's Results folder.

At the start of training, cases with 0 foreground labels are skipped; if every case has no foreground labels, training errors out. Make sure the MultiROI actually contains painted annotations.

Workflow B: Apply the Model to Predict

1. Go to the Infer tab (if you just finished training, Results folder is already filled; otherwise use the … button to select the nnUNet_results folder produced by training).

2. Click Refresh and pick the Channel to predict on in Image Channel.

3. Confirm Configuration and Fold match those used for training.

4. Optional: enter a name for the output MultiROI in Result title.

5. Click Infer -> publish MultiROI. When prediction finishes, the new MultiROI is published on the source Channel's grid and appears in the object tree.

7. Parameter Reference

About Epochs: nnU-Net v2's epoch count is fixed in its trainer class (default 1000) and has no command-line flag to change it. So the panel's Epochs field currently only scales the progress-bar display and does not alter the real training length. Set it to 0 to use nnU-Net's default.

8. Outputs

Training output. Training creates the standard nnU-Net directory layout under the Output folder: nnUNet_raw (the raw dataset, including Dataset001_Custom with imagesTr for images and labelsTr for labels — all .nii.gz — plus a dataset.json describing channels and classes), nnUNet_preprocessed (preprocessing artifacts), and nnUNet_results (the trained model). Inference must point to the nnUNet_results folder.

Inference output. On success, the predicted segmentation labels are published back into the scene as a new MultiROI overlaid on the source image Channel's grid (geometry aligned to the source Channel). The MultiROI's name comes from the Result title field (default nnUNet_prediction).

How to view. Find the new MultiROI in Dragonfly's Object Manager, toggle its visibility, and it overlays in the 2D/3D views. From there you can measure, color, and further process it like any ordinary MultiROI.

9. FAQ and Troubleshooting

Q: Setup Environment fails, reporting venv creation failure or no pip. What do I do?

A: The chosen base Python lacks a working venv/pip. Enter another CPython 3.9+ path in Base Python (build) (e.g. C:\Python312\python.exe) or a command like py -3.12, then retry. A half-built venv is rebuilt automatically.

Q: Training/inference reports CUDA not available or no GPU.

A: This plugin requires an NVIDIA GPU. Confirm a suitable NVIDIA driver is installed and that the torch CUDA build installed during Setup matches the driver (check the Torch CUDA wheel index, e.g. cu121/cu124). The smoke test at the end of Setup reports whether a usable CUDA device was found.

Q: Training reports out of memory.

A: 3D full-resolution training uses a lot of VRAM. Switch to a lighter Configuration (such as 3d_lowres or 2d), or reduce the training volume size, then retry.

Q: Clicking Train says there are no cases, or that all cases have 0 foreground labels.

A: First click Refresh, select a Channel + MultiROI pair, and click Add case (the count should be greater than 0). If it reports 0 foreground labels, the MultiROI has no painted classes — paint annotations on that MultiROI first.

Q: Dataset integrity verification fails.

A: With too few cases or unusual geometry/spacing, the check may reject the dataset. Try unticking Verify dataset integrity, or add more training cases and check that the voxel spacing is reasonable.

Q: I can't find the plugin in the menu.

A: Confirm you ticked this plugin during install (it is off by default) and that you fully restarted Dragonfly afterward (menus are scanned only at startup). If it still doesn't appear, tick it in Developer ▸ Prototype Labs... ▸ Menu Item Manager and restart again.

10. Notes and Known Limitations

• Internet + NVIDIA GPU required: the first Setup downloads several GB; training/inference need a CUDA GPU.
• Epochs field does not change training length: nnU-Net v2 has no CLI flag for epoch count, so the field only scales the progress bar (see Chapter 7).
• Turing GPUs: fp16 only: on sm_75 cards such as RTX 8000, do not use bf16.
• Training is time-consuming: 3D full-resolution training can run for hours; there is no hard timeout by default.
• Configuration/Fold must match: for inference, the chosen configuration and fold must match training, or no matching model is found.
• Single-channel images: the dataset is built with a single image channel (channel_names {"0": "image"}) and background fixed to 0.
• Windows environment: the shipped run mode is windows only; internally the plugin disables torch.compile and limits data-augmentation worker processes for Windows stability.
• Component versions come from the latest official sources at build time: nnU-Net v2's CLI arguments are version-sensitive; if an upstream update changes behavior, watch the log for error messages.

11. References

• nnU-Net (nnunetv2) official repository and docs: https://github.com/MIC-DKFZ/nnUNet
• PyTorch official site (CUDA wheel index): https://pytorch.org/ and https://download.pytorch.org/whl/
• SimpleITK official site: https://simpleitk.org/
• Full Package install / enable / uninstall instructions: the packaged README (Full Package)