nv-generate-mr-brain-finetune

NV-Generate-MR-Brain-Finetune

用途

• 用于基于用户提供的 T1、T2、FLAIR、SWI 或 MRA NIfTI 训练影像，对 NV-Generate-CTMR rflow-mr-brain v1 diffusion UNet 进行微调。
• 不用于临床解读、监管用途，也不用于批准将合成数据用于生产训练。
• 该包装器会在本地暂存配置所需文件，并将执行委托给现有的上游脚本：scripts.diff_model_create_training_data、scripts.diff_model_train，以及可选的 scripts.diff_model_infer。它不会执行 notebook。
• 清单 I/O：输入为 datalist 和 data_base_dir；输出为 finetuned_checkpoint、可选的 inference_outputs，以及 result_json。
• 底层训练契约是上游配置/环境 JSON（与 train_diff_unet_tutorial.ipynb 中 cell [10] 使用的配置相同）。包装器会为你暂存这些 JSON 文件，并通过 CLI 标志公开最常调整的字段；下文将说明这些字段、默认值，以及如何监控/调整运行。

说明

• 更改参数、副作用或验证关卡前，请先阅读 skill_manifest.yaml。
• 从 Medical AI Skills 仓库根目录运行 scripts/run_mr_brain_finetune.py。
• 如果宿主 Agent 提供 run_script，请使用 run_script("scripts/run_mr_brain_finetune.py", args=[...])；否则运行下方的 Bash/Python 命令。
• 如需检查命令格式，请勿安装软件包、克隆仓库、下载权重或启动 GPU 训练。只输出包装器命令，并包含提供的 datalist、明确指定的 --data-base-dir、明确指定的 --output-dir，以及请求的模态。
• 用户明确要求提供训练启动命令时，请勿擅自将其替换为 --preflight；只有用户请求预检时才加入 --preflight。
• 检查新的 datalist 时，请先使用 --preflight；仅当用户明确希望启动 GPU 微调时，才移除 --preflight。
• 对于暂存的预检输入包目录，如果其中存在相应文件，请将 BUNDLE/preflight_datalist.json 用作 datalist，并将 BUNDLE/preflight_dataset 用作 --data-base-dir。

示例

从输入包验证并暂存微调预检（推荐的第一步——不使用 GPU，也不进行训练）。这是唯一的标准命令；请将 INPUT_BUNDLE 和 OUT_DIR 替换为你的路径：

export NV_GENERATE_ROOT="${NV_GENERATE_ROOT:-$HOME/.cache/nvidia-skills/upstreams/NV-Generate-CTMR-da438fe}" && \
python skills/nv-generate-mr-brain-finetune/scripts/run_mr_brain_finetune.py \
  INPUT_BUNDLE/preflight_datalist.json \
  --data-base-dir INPUT_BUNDLE/preflight_dataset \
  --output-dir OUT_DIR \
  --modality mri_t1 \
  --preflight

如需实际进行 GPU 微调及其他变体，请参阅下方的 用法。

针对所请求的训练启动，检查命令格式（不进行设置或执行）：

python skills/nv-generate-mr-brain-finetune/scripts/run_mr_brain_finetune.py \
  PATH_TO_DATALIST.json \
  --data-base-dir PATH_TO_DATA_ROOT \
  --output-dir runs/nv_generate_mr_brain_finetune \
  --epochs 2 \
  --modality mri_t1

可用脚本

脚本	用途	参数
scripts/run_mr_brain_finetune.py	skill_manifest.yaml 中声明的主要入口点。	DATALIST.json --data-base-dir DATA_DIR --output-dir OUT_DIR [--epochs N] [--modality mri_t1] [--num-gpus N] [--no-amp] [--model-config FILE] [--download-model-data] [--run-inference] [--preflight]

前提条件

• 显式设置的 NV_GENERATE_ROOT 可以指向调用方的本地检出目录，并且必须包含 scripts/diff_model_create_training_data.py、scripts/diff_model_train.py 和 scripts/diff_model_infer.py。结果会记录其当前提交。
• 如果未设置 NV_GENERATE_ROOT，包装器会搜索 .workbench_data/upstreams/NV-Generate-CTMR。
• CUDA_VISIBLE_DEVICES 为可选项，可用于选择实际训练所用的 GPU。
• 运行时要求：实际训练需要 NVIDIA CUDA GPU、上游 requirements.txt 中的 Python 包，以及已下载的 MR-brain 权重。
• 副作用：在调用方提供的 --output-dir 下写入暂存配置、嵌入、检查点、可选推理图像和日志；可能会在上游检出目录和 ~/.cache/huggingface/ 下写入模型缓存；可能会访问 https://huggingface.co 获取模型资产，并访问 https://github.com 获取上游检出。
• 数据列表是 MONAI 风格的 JSON 对象，其中 training[].image 路径相对于 --data-base-dir。training[].modality 为可选项，默认值为 mri_t1。

未提供本地检出目录时，请创建一次推荐的固定版本默认检出：

if [ -z "${NV_GENERATE_ROOT:-}" ]; then
  export NV_GENERATE_COMMIT=da438fec6484cdb6f421f8c7051d954ebefff730
  export NV_GENERATE_ROOT="$HOME/.cache/nvidia-skills/upstreams/NV-Generate-CTMR-da438fe"
  if [ ! -d "$NV_GENERATE_ROOT/.git" ]; then
    git clone https://github.com/NVIDIA-Medtech/NV-Generate-CTMR.git "$NV_GENERATE_ROOT"
    git -C "$NV_GENERATE_ROOT" checkout --detach "$NV_GENERATE_COMMIT"
  fi
fi

仅当 NV_GENERATE_ROOT 位于清单指定的确切提交，且其受跟踪文件没有更改时，包装器才会执行上游代码。请通过文档说明的配置标志提供自定义训练和推理设置，而不要编辑检出内容。子进程只会接收运行时、CUDA、区域设置和证书变量的允许列表；不会转发 API 密钥、令牌、密码及其他无关的父进程环境变量。公开的 v1 资产无需凭据；如果你的网络配置需要单独的工具，请预先下载这些资产。

运行 GPU 任务前，请下载清单中声明的确切自编码器和 MR-Brain v1 检查点版本。传入 --download-model-data 会执行同样的两项固定版本下载：

python -m huggingface_hub.commands.huggingface_cli download \
  nvidia/NV-Generate-CT models/autoencoder_v1.pt \
  --revision 75ac080fb1083c403793563477724c038e7d430c \
  --local-dir "$NV_GENERATE_ROOT"
python -m huggingface_hub.commands.huggingface_cli download \
  nvidia/NV-Generate-MR-Brain models/diff_unet_3d_rflow-mr-brain_v1.pt \
  --revision ef9759bf221265b2704569cdeeac20bbf03b62ee \
  --local-dir "$NV_GENERATE_ROOT"

1. 配置和环境 JSON（根据你的数据调整）

这是对上游 train_diff_unet_tutorial.ipynb 流程的轻量封装。每次运行分为四个步骤，繁重工作交由模型作者提供的脚本完成：

1. 暂存配置 — 复制三个配置 JSON，并且只改写与本次运行相关的路径和 n_epochs（notebook 第 15 个单元格）。
2. python -m scripts.diff_model_create_training_data → 生成潜在空间 *_emb.nii.gz 嵌入（第 17 个单元格）。
3. 写入嵌入的配套文件 — 为每个嵌入创建一个 <emb>.nii.gz.json，其中包含 spacing/modality（以及模型使用时的身体区域索引）。这是 notebook 中的胶合代码（第 19 个单元格），并不位于上游 scripts/ 中，而且 diff_model_train 需要它；该技能负责这一部分。
4. python -m scripts.diff_model_train（第 21 个单元格），并可选运行 python -m scripts.diff_model_infer。

请通过编辑配置 JSON 调整参数，而不是添加标志。 所有训练/推理超参数（lr、batch_size、cache_rate、推理参数 dim/spacing/num_inference_steps/cfg_guidance_scale 等）都位于 config_maisi_diff_model_rflow-mr-brain.json 中。请编辑上游副本，或通过 --model-config FILE 提供自己的配置（另外两个配置则使用 --env-config / --model-def）。包装器只会改写下面列出的字段。

环境 JSON（environment_maisi_diff_model_rflow-mr-brain.json）——包装器会在每次运行时改写的字段：

字段	设置来源	备注
data_base_dir	--data-base-dir	相对路径 training[].image 的根目录。
json_data_list	你的数据列表	暂存副本，其中每个条目都已填入 modality。
embedding_base_dir、model_dir、output_dir	--output-dir	潜变量嵌入、检查点、推理图像。
modality_mapping_path	上游	将模态名称映射到整数代码。
model_filename	--model-filename	输出检查点名称（默认为 diff_unet_3d_rflow-mr-brain_v1.pt）。
existing_ckpt_filepath	上游权重 / --existing-ckpt-filepath	起始检查点；使用 --train-from-scratch 时会清空。
trained_autoencoder_path	上游权重 / --trained-autoencoder-path	用于对潜变量进行编码/解码的 VAE。

模型配置（config_maisi_diff_model_rflow-mr-brain.json）——封装脚本只会修改以下字段：

字段	设置来源	默认值	备注
diffusion_unet_train.n_epochs	--epochs	2（上游配置默认为 1000）	便捷覆盖项（单元格 15 也会执行相同操作）；封装脚本使用较小的默认值以便验证。
diffusion_unet_inference.modality	--modality	来自 modality_mapping.json	与训练模态保持一致，以便选择性运行 --run-inference。

该文件中的其他所有内容（lr、batch_size、cache_rate、diffusion_unet_inference 的其余字段）都保持原样——如需修改，请编辑 JSON。

固定的 v1 推理配置块默认使用 dim=[256,256,128]、 spacing=[0.94,0.94,1.36] 和 cfg_guidance_scale=2。封装脚本会保留这些字段。较旧的 v0 示例可能会使用 256^3、1 mm 间距和引导尺度 10；请以暂存的 v1 JSON 作为执行依据。

运行时标志（非配置字段）：--num-gpus N（>1 时会启动 torch.distributed.run）、--no-amp（禁用混合精度，并传递给 diff_model_train）。

--modality 从 configs/modality_mapping.json 中选择整数代码。 支持的脑部取值包括 mri（8）、mri_t1（9，默认）、mri_t2 （10）、mri_flair（11）、mri_mra（16）、mri_swi（20），以及去除颅骨的取值 mri_t1_skull_stripped（29）、 mri_t2_skull_stripped（30）、mri_flair_skull_stripped（31）、 mri_swi_skull_stripped（32）和 mri_mra_skull_stripped（33）。每个病例的 training[].modality 优先于 --modality。该模态也会传递给第 3 步的嵌入旁路文件。上游报告称，MRA 训练覆盖较少，因此无法保证 MRA 输出质量。

如需查看包含示例数据下载和检查点加载的端到端参考，请参阅上游教程 train_diff_unet_tutorial.ipynb。

2. 用法（单行训练）

仅执行预检：

export NV_GENERATE_ROOT="${NV_GENERATE_ROOT:-$HOME/.cache/nvidia-skills/upstreams/NV-Generate-CTMR-da438fe}" && \
python skills/nv-generate-mr-brain-finetune/scripts/run_mr_brain_finetune.py \
  PATH_TO_DATALIST.json \
  --data-base-dir PATH_TO_DATA_ROOT \
  --output-dir runs/nv_generate_mr_brain_finetune_preflight \
  --preflight

预检数据包输入：

export NV_GENERATE_ROOT="${NV_GENERATE_ROOT:-$HOME/.cache/nvidia-skills/upstreams/NV-Generate-CTMR-da438fe}" && \
python skills/nv-generate-mr-brain-finetune/scripts/run_mr_brain_finetune.py \
  PATH_TO_INPUT_BUNDLE/preflight_datalist.json \
  --data-base-dir PATH_TO_INPUT_BUNDLE/preflight_dataset \
  --output-dir runs/nv_generate_mr_brain_finetune_preflight \
  --preflight

GPU 微调：

export NV_GENERATE_ROOT="${NV_GENERATE_ROOT:-$HOME/.cache/nvidia-skills/upstreams/NV-Generate-CTMR-da438fe}" && \
python -m pip install -r "$NV_GENERATE_ROOT/requirements.txt" && \
python skills/nv-generate-mr-brain-finetune/scripts/run_mr_brain_finetune.py \
  PATH_TO_DATALIST.json \
  --data-base-dir PATH_TO_DATA_ROOT \
  --output-dir runs/nv_generate_mr_brain_finetune \
  --epochs 2 \
  --modality mri_t1 \
  --run-inference

将 PATH_TO_DATALIST.json 和 PATH_TO_DATA_ROOT 替换为用户的实际路径。真实训练时请勿使用 fixture 数据列表；它只是仅供预检使用的占位文件。

3. 监控训练（TensorBoard）

scripts.diff_model_train 会将 TensorBoard 事件文件写入暂存的 model_dir（OUT_DIR/artifacts/models）。使用输出目录启动 TensorBoard，并监控损失曲线：

python -m pip install tensorboard && \
tensorboard --logdir runs/nv_generate_mr_brain_finetune/artifacts

运行摘要写入 OUT_DIR/artifacts/workflow_summary.json（检查点路径、嵌入旁路文件、推理输出）；封装脚本打印到 stdout 的 JSON 会显示相同路径，另外还包含 exit_code 和用于快速排查的 stderr_tail。

4. 超参数调优和常见问题

• 损失不下降／不稳定 — 在模型配置 JSON 中调低 diffusion_unet_train.lr（默认值为 1e-5），或保持 AMP 开启（默认设置）；--no-amp 速度较慢，但在较旧的 GPU 上数值更稳定。
• 显存不足 — 在配置 JSON 中将 diffusion_unet_train.batch_size 保持为 1，并将 cache_rate 设为 0；扩容前，确认自动编码器／UNet 能在你的 GPU 上运行。多 GPU（--num-gpus N）会通过 torch.distributed.run 对批次进行分片。
• 样本较少／快速检查 — 将 --epochs 设得较小（封装脚本默认值 2 用于验证，并非用于收敛；上游配置默认提供 1000）。
• 模态条件设置错误 — 将 --modality 或每个样本的 training[].modality 设为 configs/modality_mapping.json 中存在的值；若不匹配，会明确报错，而不会静默地错误标记潜变量。
• 首次运行启动较慢 — diff_model_create_training_data 会预先计算潜在嵌入；重复使用相同的 --output-dir，以避免重新计算。

5. 评估微调后的模型

将暂存的检查点（OUT_DIR/artifacts/models/<model_filename>）用作生成时的扩散 UNet，然后检查合成的影像体：

• 在此处传入 --run-inference，快速运行内置的健全性渲染；或
• 将微调后的检查点传给 nv-generate-mr-brain 推理技能，生成新的脑部 MRI 影像体，以供定性审查。

此技能仅对文件核算和命令来源进行校验；生成图像的解剖真实性和下游效用必须由领域专家评判。

限制

• 需要使用包含现有扩散训练脚本的最新上游 NV-Generate-CTMR 检出版本。该技能会在本地暂存所需的配置和数据列表衔接内容，不依赖笔记本或 PR #33。
• 完整训练的开销可能很大，且在不同硬件、CUDA 和软件包版本下不具确定性。
• 此封装脚本只校验文件核算和命令来源，不校验解剖真实性或下游模型效用。
• 不得用于临床部署、临床解读、自主诊断、法规申报或生产环境训练数据审批。

故障排除

错误	原因	修复方法
diffusion training scripts were not found	NV_GENERATE_ROOT 未指向最新的 NV-Generate-CTMR 检出版本。	克隆或更新 https://github.com/NVIDIA-Medtech/NV-Generate-CTMR，并设置 NV_GENERATE_ROOT。
missing datalist image	training[].image 路径不是相对于 --data-base-dir 的路径，或文件不存在。	修正数据列表，或传入正确的数据根目录。
CUDA 或 MONAI 导入失败	运行环境缺少上游依赖项。	在所选环境中安装 "$NV_GENERATE_ROOT/requirements.txt"。