nv-generate-vae-finetune

NV-Generate-VAE-Finetune

用途

• 用于使用用户提供的 CT 或 MRI NIfTI 训练数据，对 NV-Generate-CTMR MAISI VAE/自动编码器进行微调。
• 不用于临床解读、监管用途，也不用于批准将合成数据用于生产训练。
• 上游目前在 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:-$HOME/.cache/nvidia-skills/upstreams/NV-Generate-CTMR-61c4ec7}" && \
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 可以指向调用方本地检出目录，且该目录 必须包含 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。

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

if [ -z "${NV_GENERATE_ROOT:-}" ]; then
  export NV_GENERATE_COMMIT=61c4ec709b84cad468852243c48e250bec732074
  export NV_GENERATE_ROOT="$HOME/.cache/nvidia-skills/upstreams/NV-Generate-CTMR-61c4ec7"
  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

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:-$HOME/.cache/nvidia-skills/upstreams/NV-Generate-CTMR-61c4ec7}" && \
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:-$HOME/.cache/nvidia-skills/upstreams/NV-Generate-CTMR-61c4ec7}" && \
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:-$HOME/.cache/nvidia-skills/upstreams/NV-Generate-CTMR-61c4ec7}" && \
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。