mcpField
T
tao-train-foundation-stereo

tao-train-foundation-stereo

github.com/nvidia開発3.8k1.6k 回インストール今日更新
リポジトリ

tao-train-foundation-stereo について

FoundationStereoによるステレオ深度推定向けのTAOスキルです。ステレオ画像ペアから視差マップを予測し、学習、評価、推論、エクスポート、量子化と、別途TensorRTデプロイのワークフローに対応します。

主な機能

  1. 01ステレオ画像ペアからの視差予測
  2. 02FoundationStereoアーキテクチャ
  3. 03学習・評価・推論
  4. 04モデルのエクスポートと量子化
  5. 05TensorRTデプロイワークフロー

活用シーン

  1. 01ステレオ画像から視差マップを生成
  2. 02対象データセットで深度モデルを学習
  3. 03TensorRTでデプロイ性能を評価

インストール

コマンドをコピーしてターミナルで実行します。

$npx skills add https://github.com/nvidia/skills --skill tao-train-foundation-stereo

インストール後、Claude Code で/tao-train-foundation-stereoを入力して読み込まれたか確認します。

同梱ファイル

SKILL.md と一緒にインストールされる参考資料やスクリプト。

BENCHMARK.md
skill-card.md
skill.oms.sig
config/
skillspector-baseline.yaml
evals/
evals.json
references/
checkpoint-inference-mappings-foundation-stereo.md
parameters-foundation-stereo.md
skill_info.yaml
spec_template_deploy.yaml
spec_template_evaluate.yaml
spec_template_export.yaml
spec_template_gen_trt_engine.yaml
spec_template_inference.yaml
spec_template_quantize.yaml
spec_template_train.yaml
spec-overrides-foundation-stereo.md
tao-deploy-foundation-stereo.md
tao-deploy-foundation-stereo.skill_info.yaml
troubleshooting-foundation-stereo.md
schemas/
evaluate.schema.json
export.schema.json
gen_trt_engine.schema.json
inference.schema.json
manifest.json
quantize.schema.json
train.schema.json

詳細

Depth Net Stereo

単独でインストールしますか? このセッションが TAO skill bank プラグインによって初期化されていない場合は、まず tao-setup スキルを実行してください(ホストの事前チェック、認証情報、スキル間の検出)。

FoundationStereo アーキテクチャを使用したステレオ深度推定です。ステレオ画像ペアから視差マップを予測し、3D 再構築を行います。

事前学習済みの Depth Anything v2 エンコーダーと EdgeNeXt エンコーダーを使用します。model.stereo_backbone.depth_anything_v2_pretrained_path と model.stereo_backbone.edgenext_pretrained_path を設定してください。

モノとステレオの両スキルは、コンテナ内で統一された TAO depth_net CLI を呼び出します。モノ/ステレオのどちらを使用するかは、model.model_type(例:FoundationStereo)で選択します。

このモデルスキルにパッケージされている PyT アクションは、train、evaluate、inference、export、quantize です。現行の TAO イメージでは、PyT の depth_net エントリーポイントは gen_trt_engine アクションを受け付けません。TensorRT エンジンのビルドには、デプロイワークフローのみを使用してください。

TAO Deploy の TensorRT アクション(gen_trt_engine、TensorRT の evaluate、TensorRT の inference)については、まず references/tao-deploy-foundation-stereo.md を読んでください。デプロイ仕様テンプレートは、このスキルの references/spec_template_deploy.yaml にあります。

Train Action Policy

このモデルはモデルレイヤーで AutoML が有効になっています。学習段階のリクエストに対応する前に、references/skill_info.yaml を読み、明示的な automl_policy の値、またはユーザーのワークフローリクエストから実行時のオーバーライドを決定してください。デフォルトでは automl_policy: on を使用し、新しい起動プロンプトでは on/off のみを提示してください。「AutoML をオフにする」「AutoML を無効にする」「HPO なし」「通常の学習」のような表現は、この実行に限り automl_policy: off として扱います。automl_policy: on、automl_enabled: true、および schemas/train.schema.json と references/spec_template_train.yaml の両方がパッケージされている場合は、デフォルトでこのモデルの skill_dir を指定して、学習アクションを tao-skill-bank:tao-run-automl 経由で実行してください。データセット、仕様、出力ディレクトリ、GPU/プラットフォーム設定、親チェックポイント、および automl_policy に関するワークフロー/アプリケーションのオーバーライドは維持してください。automl_policy: off の場合、またはパッケージされた学習スキーマ/テンプレートがない場合に限り、モデルの直接学習を使用してください。スキーマがない場合は、AutoML は有効ですが、スキーマが生成されるまでこのモデルでは実行できないことを報告してください。

evaluate、inference、export などの学習以外のアクションとデプロイフローは、このモデルスキル内にとどめてください。実行ごとの automl_policy オーバーライドによって、モデルのメタデータが変更されることはありません。

Workflow

前提条件 — データへのアクセス

データセット(左画像+右画像+GT 視差)には、コンテナ内からアクセスできる必要があります。

  • SDK runner:runner が解決する S3 パス(仕様のオーバーライドに表示される S3_TRAIN/S3_EVAL プレースホルダー)にファイルを配置してください。runner が S3 からコンテナパスへのマウントを透過的に処理します。
  • 直接の docker run(例:ローカルテスト):ホスト上のデータセットルートを、コンテナ内の同じパスに読み取り専用でマウントしてください。
docker run ... -v <host_data_root>:<host_data_root>:ro <container> ...

すべてのアクションで書き込まれる <output_dir> にも、同じアクセス要件が適用されます。

ステップ 1 — アノテーションファイル

data_sources[*].data_file で指定する、1 行ごとのアノテーションファイルです。

列数形式用途
2<left> <right>ステレオ推論(GT なし)
3<left> <right> <disparity>GT ありのステレオ
4<left> <right> <disparity> <occlusion_mask>GT とオクルージョンマスクありのステレオ

すでにアノテーションファイルがある場合は、そのパスを指定してください。ない場合は、depth_net convert で生成します。

depth_net convert -e <convert_spec.yaml>

ステレオ用の convert_spec.yaml テンプレート:

results_dir: <directory where generated annotation files are written>
data_root: <directory whose immediate children are scene folders that contain your image+depth files; convert walks data_root recursively but expects per-scene subdirectories at one level below>
image_dir_pattern: [<substring matching left image paths>]
right_dir_pattern: [<substring matching right image paths>]
depth_dir_pattern: [<substring matching GT disparity paths>]
nocc_dir_pattern: []                 # optional, occlusion mask paths
image_extension: '.png'  # always include the leading dot
depth_extension: '.png'  # form must match image_extension (the swap is a substring replace)
nocc_extension: ''
split_ratio: 0.0        # 0.0/1.0 = test-only; 0.8 = 80/20 train+val

convert は data_root を再帰的に走査し、image_dir_pattern 内のすべての部分文字列をパス文字列に含むパスを選択します(AND 条件フィルター)。その後、image_dir_pattern[0] を、対応するパターンの最初の要素に置き換え、拡張子を差し替えることで、右画像/深度/マスクのパスを導出します。データセットのディレクトリ構成を確認し、左画像、右画像、GT を区別する部分文字列(Middlebury の im0、im1、disp0GT など)を特定してください。

ステップ 2 — データに基づいて model_type と dataset_name を組み合わせます

データのレイアウトがサポート対象に合致する場合は、データセット固有のクラスを優先してください。クラス固有のパス規則、評価用クロップ、(該当する場合は)オクルージョンマスクの処理が適用されます。登録済みのクラスに合致するレイアウトがない場合に限り、GenericDataset を使用してください。

データカテゴリmodel_typedataset_name
Middlebury データFoundationStereoMiddlebury
KITTI データFoundationStereoKitti
ETH3D データFoundationStereoEth3d
FSD 合成データFoundationStereoFSD
IsaacReal 合成データFoundationStereoIsaacRealDataset
Crestereo 合成データFoundationStereoCrestereo
その他/標準外のレイアウトFoundationStereoGenericDataset

ステレオの data_sources で有効な dataset_name の値(大文字と小文字は区別されません):FSD、IsaacRealDataset、Crestereo、Middlebury、Eth3d、Kitti、GenericDataset。

同じ dataset_name の値を、学習と評価のアクション全体に適用します(いずれもGT視差を含む3列または4列のアノテーションを使用します)。デプロイ側の evaluate アクションも同じルールに従います。references/tao-deploy-foundation-stereo.md を参照してください。2列のアノテーション(左+右、GTなし)を使った推論では、データのレイアウトにかかわらず dataset_name: GenericDataset を使用してください。データセット固有のクラス(Middlebury/Kitti/Eth3d/FSD/IsaacRealDataset/Crestereo)は3列の入力を必要とし、データローダーのレベルで2列のアノテーションを拒否します。3列のアノテーション(左+右+GT)を使った推論では、データセット固有のクラスを使用できます。

ステップ 3 — spec の上書き設定から spec yaml を作成します

references/spec-overrides-foundation-stereo.md からアクションブロックをコピーします。次の項目を置き換えてください。

  • Step 2 の model.model_type(通常は FoundationStereo)
  • Step 2 の dataset.<...>.data_sources[*].dataset_name
  • Step 1 のパスを指定した dataset.<...>.data_sources[*].data_file
  • デプロイ側の evaluate では、dataset.test_dataset.batch_size: 1 を必ず設定します(references/tao-deploy-foundation-stereo.md を参照)。

形状の整合性:dataset.test_dataset.augmentation.crop_size の crop_size は、学習済みモデルの評価器とデプロイ側の TensorRT 評価器が同じ形状で動作するよう、export.input_height/input_width と一致させてください。pyt の evaluate パスでは crop_size は見た目上の設定ですが、デプロイの evaluate 側では実際に適用されます。references/troubleshooting-foundation-stereo.md と references/tao-deploy-foundation-stereo.md を参照してください。

新規インストール時のスモークテストは、crop_size: [128, 128]、dataset.max_disparity: 128、model.max_disparity: 128 で検証されています。スモークテストでは、112×112 のクロップを避け、max_disparity を正方形クロップの一辺より小さく設定しないでください。これらの組み合わせでは、チェックポイントが生成される前に FoundationStereo 内部で特徴マップまたは損失マスクの形状不一致が発生する可能性があります。

データソースの上書き設定は、すべてのアクションで必須です。各 data_sources エントリーは、必須フィールド data_file と dataset_name を持つ辞書です。アクションごとのデータセット要件表、各アクションの上書きブロック、および quantize の既知の問題に関する注記については、references/spec-overrides-foundation-stereo.md を参照してください。

ステップ 4 — 実行します

--user を使用する前に、マウントした出力パス内に書き込み可能なホーム/キャッシュディレクトリを作成してください。一部の TAO コンテナーにはホスト UID 用の /etc/passwd エントリーがなく、その UID で実行する際に PyTorch/matplotlib が書き込み可能なキャッシュパスを必要とします。

mkdir -p <output_dir>/home \
         <output_dir>/.cache/matplotlib \
         <output_dir>/.cache/torchinductor \
         <output_dir>/.cache/xdg
docker run --gpus 'device=0' --shm-size 16G --shm-size=8g \
  --user "$(id -u):$(id -g)" \
  -e USER="$(id -un)" \
  -e LOGNAME="$(id -un)" \
  -e HOME=<output_dir>/home \
  -e MPLCONFIGDIR=<output_dir>/.cache/matplotlib \
  -e TORCHINDUCTOR_CACHE_DIR=<output_dir>/.cache/torchinductor \
  -e XDG_CACHE_HOME=<output_dir>/.cache/xdg \
  -v <data_root>:<data_root>:ro \
  -v <output_dir>:<output_dir> \
  <container> \
  depth_net <action> -e <spec.yaml>

--user "$(id -u):$(id -g)" を指定しないと、コンテナーは出力を nobody:nogroup として書き込み、ホスト側でのクリーンアップ/再試行ができなくなります。

ステップ 5 — 検証
  • コンテナの終了コードが 0
  • status.json の kpi ブロックに値が設定されている
  • train の場合:各ステップの train_loss を直接確認します(エントリーポイントは、損失が NaN の場合でも Execution status: PASS を報告します)
  • evaluate の場合:epe / bp1 / bp2 / bp3 / d1 / rmse を確認します(評価器は abs_rel / sq_rel / rmse_log も出力しますが、ステレオでは意味のある値ではありません。references/parameters-foundation-stereo.md を参照してください)
  • inference の場合:results_dir 内の成果物を確認します

TAO Deploy の TensorRT アクション(gen_trt_engine、TensorRT の evaluate、TensorRT の inference)については、最初に references/tao-deploy-foundation-stereo.md を読んでください。Deploy 仕様テンプレートは、このスキルの references/ フォルダーにあり、spec_template_deploy_*.yaml というプレフィックスが付いています。

トレーニング要件

  • AutoML メトリックの契約: パッケージ化された evaluate アクションで推奨値をスコアリングする場合は、最小化方向で val/epe を使用します。
  • トレーニング監視メトリック: val/loss これは、トレーニング中の損失をリアルタイムで監視する場合にのみ使用してください。
  • AutoML/評価 KPI: val/epe を最小化します。パッケージ化された evaluate アクションはこのキーを status.json に書き込みます。AutoML の 目的関数として val/loss を使用すると、トライアルのメトリック抽出に失敗します。Lightning のコンソールテーブルでは、 同じ評価器の出力が test/epe と表示されるため、実行を検証するときは AutoML が記録した val/epe の値を、コンソールの test/epe の値と比較してください。
  • 推奨する AutoML トレーニングパラメーター: train.optim.lr と train.optim.weight_decay。検証/テスト/推論のデータ拡張フィールドは探索しないでください。また、現在の 実装では使用されないため、model.volume_dim も探索しないでください。
  • 評価データセット: 任意です。検証データセットは dataset.val_dataset.data_sources で設定します(各エントリーに data_file と dataset_name が必要です)。

アクションごとのデータセット要件をまとめた表と、各アクションで必須となるデータソースのオーバーライドブロックについては、references/spec-overrides-foundation-stereo.md を参照してください。

パラメーター、メトリクス、マルチ GPU、エクスポート/TRT、ハードウェア

重要なパラメーターの全一覧(model.encoder の vits オーバーライド、model.max_disparity のデフォルト値 416、model.volume_dim が何もしないことについての注記、dataset.baseline、dataset.focal_x、train.precision、export.batch_size を含む)、評価メトリクスの表、マルチ GPU/マルチノードの起動キー、エクスポート/TRT のデフォルト設定(opset_version/on_cpu の組み合わせ、NGC の 576×960 設定)、およびハードウェア要件については、references/parameters-foundation-stereo.md を参照してください。

エラーパターンとトラブルシューティング

視差のオーバーフロー、スモークテストの形状不一致、事前学習済みパスの欠落、encoder/dataset_name の構造体エラー、depth_net_stereo: not found というエントリーポイントの注記、pyt と deploy における crop_size の説明、deploy の evaluate で発生するスカラー変換エラーについては、references/troubleshooting-foundation-stereo.md を参照してください。

仕様パラメーター/親モデルの推論

チェックポイントの解決ルール(model_epoch_<epoch>_step_<step>.pth、dn_model_latest.pth のポリシー)、親 PyT に gen_trt_engine がないこと、depth_net_stereo.config.json に基づくアクション別推論マッピングの全一覧(parent_model/parent_job_id の解決を含む)については、references/checkpoint-inference-mappings-foundation-stereo.md を参照してください。

デプロイ

コメント