nv-generate-mr-brain-finetune

NV-Generate-MR-Brain-Finetune

目的

• ユーザーが提供する NIfTI トレーニングボリュームを使用して、NV-Generate-CTMR の rflow-mr-brain 拡散 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 です。
• 基盤となるトレーニングの契約は、上流の config/env JSON（train_diff_unet_tutorial.ipynb のセル [10] から実行するものと同じ）です。このラッパーはこれらの JSON ファイルをステージングし、最も調整頻度の高いフィールドを CLI フラグとして公開します。以下のセクションでは、各フィールド、そのデフォルト値、実行の監視方法と調整方法について説明します。

手順

• 引数、副作用、または検証ゲートを変更する前に、skill_manifest.yaml を読んでください。
• Medical AI Skills リポジトリのルートから scripts/run_mr_brain_finetune.py を実行してください。
• ホストエージェントで run_script が利用できる場合は、run_script("scripts/run_mr_brain_finetune.py", args=[...]) を使用してください。それ以外の場合は、以下の Bash/Python コマンドを実行してください。
• 新しい datalist を確認する際は、まず --preflight を使用してください。ユーザーが GPU ファインチューニングの開始を明示的に希望する場合にのみ、--preflight を削除してください。
• ステージング済みの preflight 入力バンドルディレクトリを使用する場合、該当ファイルが存在すれば、BUNDLE/preflight_datalist.json を datalist として指定し、BUNDLE/preflight_dataset を --data-base-dir に指定してください。

例

入力バンドルから preflight ファインチューニングチェックを検証してステージングします（推奨される最初の手順 — 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 は、scripts/diff_model_create_training_data.py、scripts/diff_model_train.py、および scripts/diff_model_infer.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 パッケージ、およびダウンロード済みの 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 のフローを薄くラップしたものです。各実行では、モデル作者のスクリプトに重い処理を委ねながら、次の 4 つのステップを実行します。

1. 設定ファイルを準備する — 3 つの設定 JSON をコピーし、実行ごとのパスと n_epochs のみを書き換えます（ノートブックのセル 15）。
2. python -m scripts.diff_model_create_training_data → 潜在表現の *_emb.nii.gz 埋め込みを作成します（セル 17）。
3. 埋め込みのサイドカーファイルを書き込む — 各埋め込みに対して spacing／modality（およびモデルが使用する場合は身体領域のインデックス）を含む <emb>.nii.gz.json を作成します。これは上流の scripts/ ではなく、ノートブック（セル 19）にある唯一の接続処理であり、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（他の 2 つについては --env-config／--model-def）で独自のファイルを指定してください。このラッパーが書き換えるのは、以下のフィールドだけです。

環境 JSON（environment_maisi_diff_model_rflow-mr-brain.json）— ラッパーが実行ごとに書き換えるフィールド：

フィールド	設定元	備考
data_base_dir	--data-base-dir	相対パスの training[].image のルートです。
json_data_list	ユーザーのデータリスト	各エントリーに 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 の埋め込みサイドカーファイルにも反映されます。

サンプルデータのダウンロードとチェックポイントの読み込みを含む、エンドツーエンドのリファレンスについては、上流のチュートリアル train_diff_unet_tutorial.ipynb を参照してください。

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

事前チェックのみ：

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 を、ユーザーが実際に使用するパスに置き換えてください。実際のトレーニングにはフィクスチャのデータリストを使用しないでください。これは事前チェック専用のプレースホルダーです。

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 に書き込まれます（チェックポイントのパス、埋め込みサイドカー、推論出力）。ラッパーが標準出力に出力する 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 チェックアウトが必要です。このスキル自体は必要な設定とデータリストの接続部分をローカルでステージングするため、ノートブックや 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" をインストールします。