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を実行してください。--preflightを外すのは、ユーザーがGPUでのファインチューニング開始を明示的に希望した場合に限ります。
• ステージング済みのプリフライト入力バンドルディレクトリを使用する場合は、該当するファイルがあれば、データリストにBUNDLE/preflight_datalist.jsonを、--data-base-dirにBUNDLE/preflight_datasetを指定してください。

例

入力バンドルからプリフライトのファインチューニングチェックを検証して準備します（推奨される最初の手順です。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には、configs/config_maisi_vae_train.json、scripts/transforms.py、scripts/utils.pyを含む、最新のhttps://github.com/NVIDIA-Medtech/NV-Generate-CTMRチェックアウトを指定できます。
• 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に接続する場合があります。
• データリストは、空でないtraining[]とvalidation[]またはtesting[]を含む、MONAI形式のJSONオブジェクトです。各エントリーには、--data-base-dirからの相対パスで指定したimageパスと、任意のclassまたはmodality（ctまたはmri）が含まれます。

1. 設定と環境の JSON（データに合わせて調整）

ラッパーは、上流の VAE 設定／環境 JSON を $NV_GENERATE_ROOT/configs からコピーし、以下のフィールドを書き換えて、ステージング済みのコピーを 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:-.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）

実行ツールは、イテレーションごとおよびエポックごとの 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をインストールしてください。