nv-generate-vae-finetune

NV-Generate-VAE-Finetune

用途

• 用于根据用户提供的 CT 或 MRI NIfTI 训练数据卷，微调 NV-Generate-CTMR MAISI VAE/autoencoder。
• 不可用于临床解读、监管用途，或批准将合成数据用于生产训练。
• 上游目前在 train_vae_tutorial.ipynb 中记录了 VAE 训练，并提供配置和辅助工具，但没有 scripts.train_vae CLI。本技能不会执行该 notebook；它会在本地暂存所需的配置和 datalist 连接文件，并使用上游辅助 API。
• 清单 I/O：输入为 datalist 和 data_base_dir；输出为 autoencoder_checkpoint、discriminator_checkpoint 和 result_json。
• 底层训练约定采用上游配置/环境 JSON（config_maisi_vae_train.json + environment_maisi_vae_train.json，与 train_vae_tutorial.ipynb 中的用法相同）。包装脚本会为你暂存这些 JSON 文件，并将最常调整的字段作为 CLI 标志提供；下文将介绍这些字段、默认值以及如何监控和调整运行。

说明

• 更改参数、副作用或验证门槛前，请先阅读 skill_manifest.yaml。
• 在 Medical AI Skills 仓库根目录运行 scripts/run_vae_finetune.py。
• 如果宿主 Agent 提供 run_script，请使用 run_script("scripts/run_vae_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-vae-finetune/scripts/run_vae_finetune.py \
  INPUT_BUNDLE/preflight_datalist.json \
  --data-base-dir INPUT_BUNDLE/preflight_dataset \
  --output-dir OUT_DIR \
  --modality mri \
  --preflight

真实 GPU 微调和其他用法，请参阅下文的 用法。

可用脚本

脚本	用途	参数
scripts/run_vae_finetune.py	skill_manifest.yaml 中声明的主要入口。	DATALIST.json --data-base-dir DATA_DIR --output-dir OUT_DIR [--epochs N] [--modality mri] [--patch-size 64,64,64] [--preflight]

前置条件

• NV_GENERATE_ROOT 可指向 https://github.com/NVIDIA-Medtech/NV-Generate-CTMR 的当前检出版本，其中包含 configs/config_maisi_vae_train.json、scripts/transforms.py 和 scripts/utils.py。
• 如果未设置 NV_GENERATE_ROOT，包装脚本会在 .workbench_data/upstreams/NV-Generate-CTMR 中搜索。
• CUDA_VISIBLE_DEVICES 为可选项，可用于选择实际训练所用的 GPU。
• 运行时要求：实际训练需要 NVIDIA CUDA GPU；需要上游 requirements.txt 中的 Python 包、lpips，以及下载的 VAE 权重（使用 --train-from-scratch 时除外）。
• 副作用：在调用方指定的 --output-dir 下写入暂存配置、检查点、TensorBoard 日志和运行摘要；可能会在上游检出目录、~/.cache/huggingface/ 和 ~/.cache/torch/ 下写入模型缓存；可能会访问 https://huggingface.co、https://github.com 和 https://download.pytorch.org。
• datalist 是 MONAI 风格的 JSON 对象，包含非空的 training[] 和 validation[] 或 testing[]。每个条目都有一个相对于 --data-base-dir 的 image 路径，以及可选的 class 或 modality，其值为 ct 或 mri。

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

包装器会从 $NV_GENERATE_ROOT/configs 复制上游 VAE 配置/环境 JSON，重写以下字段，并将暂存副本写入 OUT_DIR/workflow/configs/。通常只需设置数据列表和数据根目录；需要时，可使用列出的 CLI 标志覆盖单个字段。

环境 JSON（environment_maisi_vae_train.json）：

字段	设置来源	备注
model_dir	--output-dir	保存 autoencoder.pt/discriminator.pt 和最佳检查点的位置。
tfevent_path	--output-dir	TensorBoard 事件目录。
finetune	--train-from-scratch	true（默认值）会加载 trained_autoencoder_path；该标志会将其设为 false。
trained_autoencoder_path	上游权重 / --trained-autoencoder-path	微调时使用的 VAE 初始检查点。

训练字段（config_maisi_vae_train.json）：

字段	标志	类型	默认值	备注
autoencoder_train.n_epochs	--epochs	int	1
autoencoder_train.batch_size	--batch-size	int	1	每个 GPU 的批次大小（单 GPU 运行器）。
autoencoder_train.patch_size	--patch-size	int,int,int	64,64,64	训练裁剪大小。
autoencoder_train.val_batch_size	--val-batch-size	int	1
autoencoder_train.val_sliding_window_patch_size	--val-sliding-window-patch-size	int,int,int	96,96,64	滑动窗口验证 ROI。
autoencoder_train.lr	--lr	float	1e-4
autoencoder_train.perceptual_weight	--perceptual-weight	float	0.3	LPIPS 项。
autoencoder_train.kl_weight	--kl-weight	float	1e-7	KL 项。
autoencoder_train.adv_weight	--adv-weight	float	0.1	对抗项。
autoencoder_train.recon_loss	--recon-loss	l1|l2	l1
autoencoder_train.val_interval	--val-interval	int	1	两次验证之间的轮数。
autoencoder_train.cache	--cache-rate	float	0.0	MONAI CacheDataset 的缓存比例。
autoencoder_train.amp	--no-amp	flag	on	混合精度；该标志会将其禁用。
data_option.random_aug	--no-random-aug	flag	on	随机增强；该标志会将其禁用。
data_option.spacing_type	--spacing-type	original|fixed|rand_zoom	original
data_option.spacing	--spacing	float,float,float	未设置	当 spacing_type 为 fixed/rand_zoom 时必填。
data_option.select_channel	--select-channel	int	0	多通道输入使用的通道。

--modality（ct 或 mri，默认值为 mri）会为缺少 class 的数据列表条目填充该字段。必须提供验证/测试条目，因为训练循环会执行验证。

端到端参考（包括示例数据下载），请参阅上游教程 train_vae_tutorial.ipynb。

2. 用法（单行启动训练）

仅执行预检：

export NV_GENERATE_ROOT="${NV_GENERATE_ROOT:-.workbench_data/upstreams/NV-Generate-CTMR}" && \
python skills/nv-generate-vae-finetune/scripts/run_vae_finetune.py \
  PATH_TO_DATALIST.json \
  --data-base-dir PATH_TO_DATA_ROOT \
  --output-dir runs/nv_generate_vae_finetune_preflight \
  --preflight

预检捆绑包输入：

export NV_GENERATE_ROOT="${NV_GENERATE_ROOT:-.workbench_data/upstreams/NV-Generate-CTMR}" && \
python skills/nv-generate-vae-finetune/scripts/run_vae_finetune.py \
  PATH_TO_INPUT_BUNDLE/preflight_datalist.json \
  --data-base-dir PATH_TO_INPUT_BUNDLE/preflight_dataset \
  --output-dir runs/nv_generate_vae_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 -m pip install lpips tensorboard && \
python skills/nv-generate-vae-finetune/scripts/run_vae_finetune.py \
  PATH_TO_DATALIST.json \
  --data-base-dir PATH_TO_DATA_ROOT \
  --output-dir runs/nv_generate_vae_finetune \
  --epochs 1 \
  --modality mri \
  --patch-size 64,64,64 \
  --download-model-data

将 PATH_TO_DATALIST.json 和 PATH_TO_DATA_ROOT 替换为用户的实际路径。真实训练不要使用测试数据列表；它只是仅供预检使用的占位符。

3. 监控训练（TensorBoard）

运行器会在 OUT_DIR/artifacts/tfevent/autoencoder 下写入 TensorBoard 标量数据（每次迭代和每轮的 recons_loss、kl_loss、p_loss、对抗/真实/伪造损失，以及验证 scale_factor）。使用输出目录启动 TensorBoard：

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

相同的每轮损失历史也会记录在 OUT_DIR/artifacts/workflow_summary.json 中，并显示在包装器打印到 stdout 的 JSON 中（loss_history、最佳检查点路径、exit_code、stderr_tail）。

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

• 重建结果模糊 — 提高 --perceptual-weight（默认值为 0.3）；如果边缘看起来模糊不清，可尝试 --recon-loss l2。
• 后验坍塌／潜变量正则化过度 — --kl-weight 有意设为极小值（1e-7）；将其调得过高会降低重建质量。
• 对抗训练不稳定 — 降低 --adv-weight（默认值为 0.1）或 --lr；预热计划已在前 20 个 epoch 内逐步提高 LR。
• 内存不足 — 减小 --patch-size（例如 48,48,48）和 --val-sliding-window-patch-size，保持 --batch-size 1，并降低 --cache-rate。
• datalist must include non-empty validation[] or testing[] — 验证循环是必需的；请添加 validation[]（或 testing[]）条目。
• 仅支持单 GPU — 运行器会断言恰好有一个 CUDA GPU；设置 CUDA_VISIBLE_DEVICES 以选择使用哪个 GPU。

5. 评估微调后的 VAE

系统会自动跟踪验证重建损失（val_weighted_loss 最低的 epoch），并将最佳自动编码器保存为 OUT_DIR/artifacts/models 下的 autoencoder_epochN.pt。下游评估方法：

• 在 TensorBoard 中比较不同运行的验证 recons_loss/p_loss 曲线；以及
• 将微调后的自动编码器用于扩散模型微调／生成运行（例如，通过 --trained-autoencoder-path 使用 nv-generate-mr-brain-finetune），以确认潜变量仍可解码为可用的体积数据。

此技能仅把关文件核算和重建记录；图像质量和下游实用性必须由领域专家评估。

限制

• 需要使用包含 VAE 配置和辅助 API 的最新上游 NV-Generate-CTMR 检出版本。此技能负责运行器集成，不依赖 notebook。
• 完整训练成本可能很高，且不同硬件、CUDA 和软件包版本下的训练结果并不确定。
• 此封装仅把关文件核算和命令来源，不评估解剖真实性、重建质量或下游模型实用性。
• 不适用于临床部署、临床解读、自主诊断、法规申报或生产训练数据审批。

故障排除

错误	原因	解决方法
VAE configs/helpers were not found	NV_GENERATE_ROOT 未指向最新的 NV-Generate-CTMR 检出版本。	克隆或更新 https://github.com/NVIDIA-Medtech/NV-Generate-CTMR，并设置 NV_GENERATE_ROOT。
datalist must include non-empty validation[] or testing[]	VAE 训练要求为配置的验证循环提供验证数据。	添加包含相对图像路径的 validation[] 或 testing[] 条目。
CUDA、MONAI 或 LPIPS 导入失败	运行环境缺少上游依赖项。	在所选环境中安装 "$NV_GENERATE_ROOT/requirements.txt" 以及 lpips tensorboard。