nv-generate-mr-brain-finetune

NV-Generate-MR-Brain-Finetune

用途

• 使用用户提供的 NIfTI 训练体数据，对 NV-Generate-CTMR 的 rflow-mr-brain diffusion UNet 进行微调。
• 不用于临床解读、监管用途，也不用于批准将合成数据用于生产训练。
• 此包装程序会在本地暂存配置文件，并将执行工作委托给现有上游脚本：scripts.diff_model_create_training_data、scripts.diff_model_train，以及可选的 scripts.diff_model_infer。它不会执行笔记本。
• 清单输入/输出：输入为 datalist 和 data_base_dir；输出为 finetuned_checkpoint、可选的 inference_outputs，以及 result_json。
• 底层训练约定采用上游配置/环境变量 JSON（与 train_diff_unet_tutorial.ipynb 中单元格 [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 命令。
• 检查新的 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:-.workbench_data/upstreams/NV-Generate-CTMR}" && \
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 微调及其他用法，请参阅下方的用法。

可用脚本

脚本	用途	参数
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] [--run-inference] [--preflight]

先决条件

• NV_GENERATE_ROOT 可以指向当前检出的 https://github.com/NVIDIA-Medtech/NV-Generate-CTMR，其中包含 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 获取上游检出内容。
• datalist 是一个 MONAI 风格的 JSON 对象，其中 training[].image 路径相对于 --data-base-dir。training[].modality 为可选项，默认值为 mri_t1。

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. 写入嵌入 sidecar 文件 — 为每个嵌入创建一个 <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	你的 datalist	暂存的副本，并为每个条目填入 modality。
embedding_base_dir、model_dir、output_dir	--output-dir	潜在嵌入、检查点、推理图像。
modality_mapping_path	上游	将模态名称映射为整数代码。
model_filename	--model-filename	输出检查点名称（默认值为 diff_unet_3d_rflow-mr-brain_v0.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 即可更改这些内容。

运行时标志（非配置字段）：--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_swi（20），以及对应的 *_skull_stripped 变体（29/30/31/32）。每个病例中的 training[].modality 会覆盖 --modality。该模态也会用于第 3 步的嵌入 sidecar 文件。

端到端参考（包括示例数据下载和检查点加载），请参阅上游教程 train_diff_unet_tutorial.ipynb。

2. 用法（单行训练）

仅执行预检：

export NV_GENERATE_ROOT="${NV_GENERATE_ROOT:-.workbench_data/upstreams/NV-Generate-CTMR}" && \
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:-.workbench_data/upstreams/NV-Generate-CTMR}" && \
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:-.workbench_data/upstreams/NV-Generate-CTMR}" && \
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 datalist；它只是用于预检的占位文件。

3. 监控训练（TensorBoard）

scripts.diff_model_train 会在暂存的 model_dir（OUT_DIR/artifacts/models）下写入 TensorBoard 事件文件。针对输出目录启动 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 检出版本，其中包含现有的扩散训练脚本。此技能本身会在本地暂存所需的配置和数据列表衔接内容，不依赖 notebook 或 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"。