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は提供していません。このスキルはノートブックを実行せず、必要な設定ファイルとデータリストの連携ファイルをローカルに用意して、アップストリームのヘルパーAPIを使用します。
• マニフェストの入出力：入力は datalist と data_base_dir、出力は autoencoder_checkpoint、discriminator_checkpoint、result_json です。
• 基盤となるトレーニングの仕様は、アップストリームの設定／環境JSON（train_vae_tutorial.ipynb で使用される config_maisi_vae_train.json と environment_maisi_vae_train.json）です。ラッパーがこれらのJSONファイルを用意し、最も調整される項目をCLIフラグとして公開します。以下のセクションでは、各項目とそのデフォルト値、および実行の監視・調整方法について説明します。

手順

• 引数、副作用、検証ゲートを変更する前に、skill_manifest.yaml を読んでください。
• Medical AI Skills リポジトリのルートから scripts/run_vae_finetune.py を実行してください。
• ホストエージェントが run_script を提供する場合は、run_script("scripts/run_vae_finetune.py", args=[...]) を使用してください。そうでない場合は、以下のBash／Pythonコマンドを実行してください。
• 新しいデータリストを確認する場合は、まず --preflight を使用してください。GPUでのファインチューニングを開始するようユーザーが明示的に希望した場合に限り、--preflight を削除してください。
• ステージング済みの入力バンドルディレクトリを使う場合は、該当するファイルがあれば、BUNDLE/preflight_datalist.json をデータリストとして、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 を使用する場合はVAE重みは不要です。
• 副作用：呼び出し元が指定する --output-dir の下に、ステージング済み設定ファイル、チェックポイント、TensorBoardログ、実行サマリーを書き込みます。アップストリームのチェックアウト、~/.cache/huggingface/、~/.cache/torch/ の下にモデルキャッシュを書き込む場合があります。また、https://huggingface.co、https://github.com、https://download.pytorch.org に接続する場合があります。
• データリストは、空でない training[] と validation[] または testing[] を含むMONAI形式のJSONオブジェクトです。各エントリには --data-base-dir からの相対パスで指定した image パスと、任意の class または ct か mri の modality が含まれます。

ローカルのチェックアウトが指定されていない場合は、推奨される固定バージョンのチェックアウトを一度作成してください。

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 が設定されていないデータリスト項目に対して、項目ごとの class を設定します。トレーニングループでは検証パスが実行されるため、検証／テスト項目が必要です。

サンプルデータのダウンロードを含むエンドツーエンドのリファレンスについては、上流のチュートリアル train_vae_tutorial.ipynb を参照してください。

2. 使用方法（1 行でトレーニング）

事前チェックのみ：

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）

ランナーは、TensorBoard のスカラー値（イテレーションごとおよびエポックごとの recons_loss、kl_loss、p_loss、敵対的／実／偽の損失、および検証用の scale_factor）を OUT_DIR/artifacts/tfevent/autoencoder に書き込みます。出力ディレクトリを指定して TensorBoard を起動します：

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

同じエポックごとの損失履歴は OUT_DIR/artifacts/workflow_summary.json にも記録され、ラッパーが標準出力に出力する 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エポックにわたって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がちょうど1つであることを確認します。使用するGPUを選択するには、CUDA_VISIBLE_DEVICESを設定してください。

5. ファインチューニング済みVAEを評価する

検証再構成損失（val_weighted_lossが最小のエポック）は自動的に追跡され、最良のオートエンコーダーは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チェックアウトが必要です。このスキルはランナーの接続処理を担い、ノートブックには依存しません。
• 学習全体には多大なコストがかかる場合があり、ハードウェア、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をインストールしてください。