MRが楽しい

MRやVRについて学習したことを書き残す

Unity AIのドキュメントを読む その65(テンソルデータへの直接アクセス)

本日はUnityの技術調査枠です。
Unity AIのドキュメントを読みながら実際に操作を試して記事に残します。

Unity AI

以下のUnity AIのドキュメントを試しながら実行時のキャプチャをしていきます。
docs.unity3d.com

テンソルデータへの直接アクセス

テンソルにアクセスしたり、モデル間でテンソルを渡したりする際に、遅いリードバックが発生島sう。
これを回避するにはテンソルの基盤となるネイティブデータに対して直接読み取りと書き込みを行います。

テンソルデータの保存場所を確認する

dataOnBackend.backendTypeプロパティを使用してテンソルデータが以下の場所に保存されているかどうかを確認します。

  • BackendType.CPU
  • BackendType.GPUCompute
  • BackendType.GPUPixel
using UnityEngine;
using Unity.InferenceEngine;

public class CheckTensorLocation : MonoBehaviour
{
    public Texture2D inputTexture;

    void Start()
    {
        // 入力データをテンソルとして作成する
        Tensor<float> inputTensor = new Tensor<float>(new TensorShape(1, 4, inputTexture.height, inputTexture.width));
        TextureConverter.ToTensor(inputTexture, inputTensor);

        // テンソルがCPUまたはGPUメモリに保存されているかどうかを確認し、コンソールウィンドウに書き込みます
        Debug.Log(inputTensor.dataOnBackend.backendType);
    }
}

テンソルを他のデバイスに強制的に配置したい場合は以下を使用します。

  • ComputeTensorData.PinはテンソルをComputeBuffer内のGPUのコンピュートシェーダーメモリに強制的に配置します。
  • CPUTensorData.PinはテンソルをCPUメモリに強制的に配置します。
// テンソルを作成する
Tensor<float> inputTensor = new Tensor<float>(new TensorShape(1, 3, 2, 2));

// テンソルをGPUメモリに強制的に保存する
ComputeTensorData computeTensorData = ComputeTensorData.Pin(inputTensor);
Note

テンソルデータがターゲットデバイスに既に存在する場合、このメソッドはそれを変更せずにそのまま渡します。
そうでない場合は既存のデータを破棄し、選択されたバックエンドに新しいメモリを割り当てます。

CPUデータに直接アクセスする

テンソルがCPU上にあり、それに依存するすべての演算が完了するとテンソルは読み取りと書き込みが可能になります。
インデクサーを使用してテンソルデータを操作できます。

var tensor = new Tensor<float>(new TensorShape(1, 2, 3));
//...
if (tensor.backendType == BackendType.CPU && tensor.IsReadbackRequestDone()) {
    // テンソルは直接読み書き可能な状態です
    tensor[0, 1, 0] = 1f;
    tensor[0, 1, 1] = 2f;
    tensor[0, 1, 2] = 3f;
    float val = tensor[0, 0, 2];
}

テンソルの読み取り可能なフラット化バージョンをspanまたはNativeArrayとして取得することもできます。
データはテンソルの行優先のフラット化されたメモリレイアウトになります。

var tensor = new Tensor<float>(new TensorShape(1, 2, 3));
//...
if (tensor.backendType == BackendType.CPU && tensor.IsReadbackRequestDone()) {
    // テンソルは読み取り可能な状態です
    var nativeArray = tensor.AsReadOnlyNativeArray();
    float val010 = nativeArray[3 + 0];
    float val011 = nativeArray[3 + 1];
    float val012 = nativeArray[3 + 2];

    var span = tensor.AsReadOnlySpan();
    float val002 = span[2];
}

バックエンドメモリにデータを直接アップロード

Uploadを使用してデータをテンソルに直接アップロードします。

var tensor = new Tensor<float>(new TensorShape(1,2,3), new [] { 0f, 1f, 2f, 3f, 4f, 5f });
tensor.Upload(new [] { 6f, 7f, 8f });
// tensor dataOnBackend now contains {6,7,8,3,4,5}

このメソッドは全てのテンソルデータバックエンドで機能しますがブロッキング呼び出しになる可能性があります。
テンソルデータがCPU上にある場合、推論エンジンはテンソルの保留中のジョブが完了するまでそれをブロックします。
テンソルデータがGPU上にある場合、推論エンジンはGPUアップロードを実行します。

GPUメモリ内のテンソルへのアクセス

GPUコンピューティングメモリに格納されているテンソルにアクセスするにはComputeTensorData.Pinを使用して、データをComputeTensorDataオブジェクトとして取得します。
その後bufferプロパティを使用して、コンピューティングバッファ内のテンソルデータに直接アクセスできます。

CPUメモリ内のテンソルへのアクセス

CPUメモリに格納されているテンソルにアクセスするには、CPUTensorData.Pin を使用して、データを CPUTensorData オブジェクトとして取得します。
このオブジェクトはIJobParallelForなどのBurst関数で使用して、テンソルデータの読み取りと書き込みを行うことができます。
オブジェクトの読み取りフェンス(CPUTensorData.fence)と書き込みフェンス(CPUTensorData.reuse)プロパティを使用して、Burstジョブの依存関係を処理します。

NativeTensorArrayクラスのメソッドを使用して、テンソルデータをネイティブ配列として読み取りと書き込みを行うこともできます。