Unity AIのドキュメントを読む その６０（モデルからの出力を非同期に読み取る）

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

Unity AI

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

モデルからの出力を非同期に読み取る

モデルをスケジュールし、PeekOutputから出力テンソルにアクセスしたときに以下の状況が発生します。

• 推論エンジンが最終テンソルデータの計算を完了していない可能性があるため、スケジュール済みの作業が保留中となる。
• グラフィックス処理装置 (GPU) バックエンドを使用している場合、計算されたテンソルデータはGPU上にある可能性がある。

これらの条件のいずれかが当てはまる場合、ReadbackAndCloneメソッドまたはCompleteAllPendingOperationsメソッドは操作が完了するまでメインスレッドをブロックします。
これを回避するには以下の方法に従って非同期リードバックを使用します。

1. 待機可能なReadbackAndCloneAsyncメソッドを使用します。
推論エンジンは入力テンソルのCPUコピーを非ブロッキング方式で返します。

using Unity.InferenceEngine;
using UnityEngine;

public class AsyncReadbackCompute : MonoBehaviour
{
    [SerializeField]
    ModelAsset modelAsset;

    Tensor<float> m_Input;
    Worker m_Worker;

    async void OnEnable()
    {
        var model = ModelLoader.Load(modelAsset);
        m_Input = new Tensor<float>(new TensorShape(1, 1), new[] { 43.0f });
        m_Worker = new Worker(model, BackendType.GPUCompute);
        m_Worker.Schedule(m_Input);

        // テンソルの所有権を取得せずに推論エンジンから値を覗き見る
        var outputTensor = m_Worker.PeekOutput() as Tensor<float>;
        var cpuCopyTensor = await outputTensor.ReadbackAndCloneAsync();

        Debug.Assert(cpuCopyTensor[0] == 42);
        Debug.Log($"Output tensor value {cpuCopyTensor[0]}");
        cpuCopyTensor.Dispose();
    }

    void OnDisable()
    {
        m_Input.Dispose();
        m_Worker.Dispose();
    }
}

2. ReadbackRequestメソッドとTensor.IsReadbackRequestDoneメソッドでポーリングメカニズムを使用します。

bool inferencePending = false;
Tensor<float> outputTensor;

void OnUpdate()
{
    if (!inferencePending)
    {
        m_Worker.Schedule(m_Input);
        outputTensor = m_Worker.PeekOutput() as Tensor<float>;

        // Trigger a non-blocking readback request
        outputTensor.ReadbackRequest();
        inferencePending = true;
    }
    else if (inferencePending && outputTensor.IsReadbackRequestDone())
    {
        // m_OutputがCPUにダウンロードされました
        // ReadbackAndCloneまたはToReadOnlyArrayを使用してもブロックされません
        var array = outputTensor.DownloadToArray();
        inferencePending = false;
    }
}

3. コールバック付きのawaitableを使用します。

bool inferencePending = false;

void Update()
{
    if (!inferencePending)
    {
        m_Worker.Schedule(m_Input);
        var outputTensor = m_Worker.PeekOutput() as Tensor<float>;
        inferencePending = true;

        var awaiter = outputTensor.ReadbackAndCloneAsync().GetAwaiter();
        awaiter.OnCompleted(() =>
        {
            var tensorOut = awaiter.GetResult();
            inferencePending = false;
            tensorOut.Dispose();
        });
    }
}