NVIDIAは、Jetson AGX Orinのディープラーニングアクセラレーター（DLA）で推論を実行・管理するPythonバインディングを公開した。開発者はC++でランタイムを制御せずに推論をDLAに任せ、同じプログラムでGPU処理と接続できる。

1. ランタイム管理をPythonで行う

新しいバインディングは、CUDA Pythonプロジェクトのcuda-bindings 13.3.0に含まれるcuda.bindings.cudlaだ。バインディングは既存のAPIを別の言語から使えるようにつなぐ層で、今回はデバイスの作成から推論タスクの投入までPythonで処理できる。

プログラムの実行を支える基盤ソフトウェアであるDLAランタイムは、JetPackに含まれるため別途インストールする必要はない。コンパイル済みの、実行可能なコードに変換されたモデルの実行にはこのランタイムを使い、TensorRTはモデルを実行可能な形式に変換するときだけ必要だ。

バインディングが要求するPythonの最小バージョンは3.10だ。NVIDIAは、必要に応じて3.12の仮想環境を用意してインストールする方法も案内している。

2. モデル変換とデバイス上での実行を分ける

DLA loadableは、学習済みネットワークを専用アクセラレーターが直接読み込めるように変換したバイナリーファイルだ。TensorRTのtrtexecで一度作成すれば、推論のたびにモデル変換ツールを実行し直す必要はない。

NVIDIAのビルド例では、Jetson AGX OrinでJetPack 7.2.2以降を使用している。

ツールのバージョンはTensorRT 11.4.1とCUDA 13.4 Update 1だ。

以下のコマンドは、入力サイズが固定されたONNXモデルを対象とする。モデルのテンソルは16ビット浮動小数点形式で、コマンドでは推論演算を行う実行単位である最初のDLAエンジンと、データの並び順を表す入出力データのレイアウトを指定する。

trtexec \
    --onnx=your_model_fp16.onnx \
    --useDLACore=0 \
    --inputIOFormats=chw16 \
    --outputIOFormats=chw16 \
    --memPoolSize=dlaSRAM:1 \
    --saveEngine=model.bin \
    --buildDLAStandalone

このコマンドで作成するファイルは、推論時にモデル変換ツールなしで実行するDLA用バイナリーだ。ビルド成果物と実行を分離すれば、Pythonアプリケーションは準備済みのモデルを読み込み、タスクを投入することに集中できる。

3. 1つのタスクキューでGPUとDLAをつなぐ

GPUとDLAが連携するハイブリッドモードでは、タスクを順番に投入するキューであるCUDAストリームを使う。GPUでの前処理、DLA推論、GPUでの後処理を同じキューに入れると、入力の準備から結果の処理まで1つの流れにつながる。

実行は、CUDAを初期化してモデルを読み込んだ後、多次元に配列された数値データであるモデルのテンソル情報に合わせたGPUメモリーをDLAに登録する順序で始まる。入力をコピーして推論タスクを投入し、完了した結果をCPUメモリーに取り込んでPythonで後処理する。

モデルから読み取ったテンソル情報には、データを一時的に保持するメモリー領域であるバッファーのサイズとデータ形式が含まれる。これをメモリー割り当てと入出力データの配置に活用すれば、モデルごとにサイズをコードに固定する手間を減らせる。返されたバイト配列もこの情報に合わせて解釈し、Pythonの後続処理に渡す。

4. GPUの余力を得る代わりにモデルとメモリーを管理する

NVIDIAの説明によると、対応する演算をDLAに任せれば、電力効率を高め、GPUの余力をアプリケーションのほかの処理に残せる。開発者はアクセラレーターを使い分けながら、実行を調整するコードをPython内に維持できる。

DLAはGPUより対応レイヤーが少なく、モデル内で演算が接続された構造であるモデルグラフの前後の演算を調整する必要が生じる場合がある。開発者はバッファーの登録・コピー、タスクの完了時点を合わせる同期処理、リソースの解放を担う必要がある。今回のバインディングはハイブリッドモードのみに対応し、デバイスハンドルはPythonの自動メモリー回収では解放されない。

5. エラー処理と終了もPythonの流れにまとめる

cuDLAの呼び出しに失敗すると、ステータスコードと読み取り可能なエラー名を含むPythonの例外が発生する。開発者はこれを受け取り、不正な引数やハードウェアエラーに応じた処理を記述できる。

NVIDIAは、クリーンアップコードをtry/finallyに置く方法を推奨している。登録メモリー、モデルモジュールとデバイス、GPUメモリー、ストリームを、それぞれの作成に対応する解放呼び出しで片付ける手順だ。モデルの実行だけでなく、エラー対応と終了まで1つのプログラムの流れにまとめることになる。

cuDLAのPythonバインディングは、専用アクセラレーターに推論を任せ、GPUをほかの処理に活用したいエッジ開発者に、実行フローを1つの言語で管理する選択肢を提供する。