nv-generate-mr-brain-finetune

NV-Generate-MR-Brain-Finetune

目的

• ユーザーが提供した T1、T2、FLAIR、SWI、または MRA の NIfTI トレーニングボリュームを使用して、NV-Generate-CTMR rflow-mr-brain v1 拡散 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 です。
• 基盤となるトレーニングの仕様は、上流の設定／環境変数 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 コマンドを実行してください。
• コマンド形式の確認では、パッケージのインストール、リポジトリのクローン、重みのダウンロード、GPU トレーニングの開始は行わないでください。指定された datalist、明示的な --data-base-dir、明示的な --output-dir、および要求されたモダリティを含むラッパーコマンドのみを出力してください。
• ユーザーがトレーニング開始コマンドを明示的に求めた場合、黙って --preflight に置き換えないでください。事前確認の要求に限り --preflight を含めてください。
• 新しい datalist を確認する際は、まず --preflight を使用してください。ユーザーが GPU ファインチューニングの開始を明示的に希望した場合に限り、--preflight を外してください。
• ステージング済みの事前確認用入力バンドルディレクトリでは、該当するファイルが存在する場合、datalist に BUNDLE/preflight_datalist.json を、--data-base-dir に BUNDLE/preflight_dataset を使用してください。

例

入力バンドルから、事前確認用のファインチューニングチェックを検証してステージングします（推奨される最初のステップです。GPU もトレーニングも不要です）。これが唯一の標準コマンドです。INPUT_BUNDLE と OUT_DIR をパスに置き換えてください。

export NV_GENERATE_ROOT="${NV_GENERATE_ROOT:-$HOME/.cache/nvidia-skills/upstreams/NV-Generate-CTMR-da438fe}" && \
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 ファインチューニングやその他のバリエーションについては、以下の使用方法を参照してください。

要求されたトレーニング開始コマンドの形式確認（セットアップや実行は行いません）。

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

利用可能なスクリプト

スクリプト	目的	引数
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] [--download-model-data] [--run-inference] [--preflight]

前提条件

• 明示的な NV_GENERATE_ROOT には、呼び出し元のローカルチェックアウトを指定できます。このディレクトリには scripts/diff_model_create_training_data.py、scripts/diff_model_train.py、および scripts/diff_model_infer.py が含まれている必要があります。結果には現在のコミットが記録されます。
• 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 に接続する場合があります。
• データリストは MONAI 形式の JSON オブジェクトで、training[].image のパスは --data-base-dir からの相対パスです。training[].modality は任意で、省略すると mri_t1 が使用されます。

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

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

ラッパーは、NV_GENERATE_ROOT がマニフェストで指定された正確なコミットにあり、追跡対象ファイルに変更がない場合にのみ、アップストリームのコードを実行します。カスタムのトレーニング設定と推論設定は、チェックアウトを編集せず、ドキュメントに記載された設定フラグで指定してください。子プロセスには、ランタイム、CUDA、ロケール、および証明書に関する環境変数の許可リストのみが渡されます。API キー、トークン、パスワード、および無関係な親プロセスの環境変数は転送されません。公開されている v1 アセットに認証情報は不要です。ネットワーク環境に別途ツールが必要な場合は、事前にダウンロードしてください。

GPU を使って実行する前に、マニフェストで宣言された正確なオートエンコーダーと MR-Brain v1 チェックポイントのリビジョンをダウンロードしてください。--download-model-data を指定すると、同じ 2 つの固定バージョンのダウンロードが実行されます。

python -m huggingface_hub.commands.huggingface_cli download \
  nvidia/NV-Generate-CT models/autoencoder_v1.pt \
  --revision 75ac080fb1083c403793563477724c038e7d430c \
  --local-dir "$NV_GENERATE_ROOT"
python -m huggingface_hub.commands.huggingface_cli download \
  nvidia/NV-Generate-MR-Brain models/diff_unet_3d_rflow-mr-brain_v1.pt \
  --revision ef9759bf221265b2704569cdeeac20bbf03b62ee \
  --local-dir "$NV_GENERATE_ROOT"

1. 設定 JSON と環境 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_v1.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 を編集してください。

固定された v1 推論ブロックのデフォルト値は、dim=[256,256,128]、 spacing=[0.94,0.94,1.36]、cfg_guidance_scale=2 です。ラッパーは これらのフィールドをそのまま保持します。古い v0 の例では、256^3、1 mm の間隔、ガイダンス スケール 10 が示されている場合があります。実行時の正しい設定として、ステージング済みの v1 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_mra（16）、mri_swi（20）と、 頭蓋骨除去済みの値である mri_t1_skull_stripped（29）、 mri_t2_skull_stripped（30）、mri_flair_skull_stripped（31）、 mri_swi_skull_stripped（32）、mri_mra_skull_stripped（33）があります。ケースごとの training[].modality は --modality より優先されます。モダリティは ステップ 3 の埋め込みサイドカーにも使用されます。上流の報告によると、MRA の学習データは少ないため、 MRA の出力品質は保証されません。

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

2. 使用方法（1 行で学習）

事前チェックのみ：

export NV_GENERATE_ROOT="${NV_GENERATE_ROOT:-$HOME/.cache/nvidia-skills/upstreams/NV-Generate-CTMR-da438fe}" && \
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:-$HOME/.cache/nvidia-skills/upstreams/NV-Generate-CTMR-da438fe}" && \
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:-$HOME/.cache/nvidia-skills/upstreams/NV-Generate-CTMR-da438fe}" && \
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" をインストールしてください。