UE5.4 での NNE Runtime - ほげたつブログ

生きてます。ONNX Runtime を C++ で組み込んだよって話を最近聞いたので UE の NNE Runtime を試そうとしたら、5.4 で微妙に変わってたようなので備忘録。

…と、この記事を書き始めてから見つけてしまったのですが、大変詳しく書かれている以下の資料があったので UE5.4 における NNE Runtime についてはこれを読めば良さそうです。この記事では備忘録としてサンプルコードだけ残しておきます。

www.docswell.com

サンプルコード

みんな大好き MNIST の ONNX モデルを使って UE の C++ から呼び出して結果を受け取るところまでのサンプルです。MNIST は 28x28 画像を受け取って 0-9 の数字のどれが一番近そうか推論するモデルですね。この動画では 28x28 の R32F 画像をリストで持ち、その中から一つランダムで選んで NNE Runtime 経由で ONNX モデルに入力として渡した後に結果を受け取っています。

UE5.4になってNNE Runtimeのプラグイン構成や使い方が微妙に変わってたので試してみてた。後で備忘録としてまとめる pic.twitter.com/zgM0vtHdK4
— ほげたつ (@HogeTatu) 2024年6月30日

ざっくり処理の流れは以下です。

SetupNNE() でモデルの読み込みとインスタンスの作成
RunNNE() をブループリントから実行
- モデルに入力として渡すテクスチャを決定
- OnStartNNE(const UTexture2D* InputTexture) 呼び出し
- AsyncTask を使い AnyThread から NNE Runtime を実行
- 出力を受け取ったら GameThread で OnEndNNE(const TArray& OutputData) 呼び出し

ヘッダー

#pragma once

#include "CoreMinimal.h"
#include "GameFramework/Actor.h"
#include "TestActor.generated.h"

class UTexture2D;
class UNNEModelData;

namespace UE::NNE
{
    class IModelCPU;
    class IModelInstanceCPU;
    struct FTensorBindingCPU;
}

class FNNEModelHelper
{
public:
    TSharedPtr<UE::NNE::IModelInstanceCPU> ModelInstance;
    TWeakObjectPtr<UTexture2D> InputTexture;
    TArray<float> OutputData;
    TArray<UE::NNE::FTensorBindingCPU> InputBindings;
    TArray<UE::NNE::FTensorBindingCPU> OutputBindings;
    bool bIsRunning = false;
};

UCLASS()
class ATestActor : public AActor
{
    GENERATED_BODY()

public:
    UPROPERTY(EditAnywhere, meta = (DisplayName = "NNE Model Data"))
    TObjectPtr<UNNEModelData> NNEModelData;

    UPROPERTY(EditAnywhere)
    TArray<TObjectPtr<UTexture2D>> InputTextures;

private:
    TSharedPtr<UE::NNE::IModelCPU> Model;
    TSharedPtr<UE::NNE::IModelInstanceCPU> ModelInstance;
    TSharedPtr<FNNEModelHelper> ModelHelper;
    
public:    
    ATestActor();

protected:
    virtual void BeginPlay() override;
    virtual void EndPlay(const EEndPlayReason::Type EndPlayReason) override;

private:
    void SetupNNE();
    void DestroyNNE();

public:
    // NNE実行
    UFUNCTION(BlueprintCallable, Category = "TestActor")
    void RunNNE();

    // NNE実行開始通知
    UFUNCTION(BlueprintImplementableEvent, Category = "TestActor")
    void OnStartRunNNE(const UTexture2D* InputTexture);

    // NNE実行終了通知
    UFUNCTION(BlueprintImplementableEvent, Category = "TestActor")
    void OnEndRunNNE(const TArray<float>& OutputData);
};

ソース

#include "TestActor.h"
#include "NNE.h"
#include "NNEModelData.h"
#include "NNERuntimeCPU.h"
#include "Engine/Texture2D.h"

ATestActor::ATestActor()
{
    PrimaryActorTick.bCanEverTick = true;
}

void ATestActor::BeginPlay()
{
    Super::BeginPlay();
    SetupNNE();
}

void ATestActor::EndPlay(const EEndPlayReason::Type EndPlayReason)
{
    Super::EndPlay(EndPlayReason);
    DestroyNNE();
}

void ATestActor::SetupNNE()
{
    TWeakInterfacePtr<INNERuntimeCPU> Runtime = UE::NNE::GetRuntime<INNERuntimeCPU>(FString("NNERuntimeORTCpu"));
    if (!Runtime.IsValid())
    {
        UE_LOG(LogTemp, Error, TEXT("Cannot find runtime NNERuntimeORTCpu, please enable the corresponding plugin"));
    }

    Model = Runtime->CreateModelCPU(NNEModelData);
    if (!Model.IsValid())
    {
        UE_LOG(LogTemp, Error, TEXT("Failed to create the model"));
    }

    ModelInstance = Model->CreateModelInstanceCPU();
    if (!ModelInstance.IsValid())
    {
        UE_LOG(LogTemp, Error, TEXT("Failed to create the model instance"));
    }

    ModelHelper = MakeShared<FNNEModelHelper>();
}

void ATestActor::DestroyNNE()
{
    if (ModelHelper.IsValid())
    {
        ModelHelper.Reset();
    }

    if (ModelInstance.IsValid())
    {
        ModelInstance.Reset();
    }

    if (Model.IsValid())
    {
        Model.Reset();
    }
}

void ATestActor::RunNNE()
{
    if (!ModelInstance.IsValid())
    {
        UE_LOG(LogTemp, Error, TEXT("Model instance is not valid"));
        return;
    }
    if (InputTextures.Num() == 0)
    {
        UE_LOG(LogTemp, Error, TEXT("Input textures is empty"));
        return;
    }
    if (ModelHelper->bIsRunning)
    {
        UE_LOG(LogTemp, Error, TEXT("Model is already running"));
        return;
    }

    ModelHelper->InputTexture = InputTextures[FMath::RandRange(0, InputTextures.Num() - 1)];
    if (!ModelHelper->InputTexture.IsValid())
    {
        UE_LOG(LogTemp, Error, TEXT("Input texture is not valid"));
        return;
    }
    if (ModelHelper->InputTexture->GetPixelFormat() != PF_R32_FLOAT)
    {
        UE_LOG(LogTemp, Error, TEXT("Input texture format is not supported"));
        return;
    }

    ModelHelper->ModelInstance = ModelInstance;
    ModelHelper->bIsRunning = true;

    OnStartRunNNE(ModelHelper->InputTexture.Get());

    TWeakObjectPtr<ATestActor> WeakThis = this;
    AsyncTask(ENamedThreads::AnyNormalThreadNormalTask, [WeakThis, ModelHelper = ModelHelper]()
    {
        TConstArrayView<UE::NNE::FTensorDesc> InputTensorDescs = ModelHelper->ModelInstance->GetInputTensorDescs();
        checkf(InputTensorDescs.Num() == 1, TEXT("The current example supports only models with a single input tensor"));
        UE::NNE::FSymbolicTensorShape SymbolicInputTensorShape = InputTensorDescs[0].GetShape();
        checkf(SymbolicInputTensorShape.IsConcrete(), TEXT("The current example supports only models without variable input tensor dimensions"));
        TArray<UE::NNE::FTensorShape> InputTensorShapes = { UE::NNE::FTensorShape::MakeFromSymbolic(SymbolicInputTensorShape) };

        ModelHelper->ModelInstance->SetInputTensorShapes(InputTensorShapes);

        TConstArrayView<UE::NNE::FTensorDesc> OutputTensorDescs = ModelHelper->ModelInstance->GetOutputTensorDescs();
        checkf(OutputTensorDescs.Num() == 1, TEXT("The current example supports only models with a single output tensor"));
        UE::NNE::FSymbolicTensorShape SymbolicOutputTensorShape = OutputTensorDescs[0].GetShape();
        checkf(SymbolicOutputTensorShape.IsConcrete(), TEXT("The current example supports only models without variable output tensor dimensions"));

        ModelHelper->InputBindings.SetNumZeroed(1);
        ModelHelper->OutputBindings.SetNumZeroed(1);

        // Fill the input tensor with the input texture data
        FTexture2DMipMap& Mip = ModelHelper->InputTexture->GetPlatformData()->Mips[0];
        ModelHelper->InputBindings[0].Data = Mip.BulkData.Lock(LOCK_READ_ONLY);
        ModelHelper->InputBindings[0].SizeInBytes = Mip.BulkData.GetBulkDataSize();

        // Allocate memory for the output tensor
        ModelHelper->OutputData.SetNumUninitialized(10);
        ModelHelper->OutputBindings[0].Data = ModelHelper->OutputData.GetData();
        ModelHelper->OutputBindings[0].SizeInBytes = ModelHelper->OutputData.Num() * sizeof(float);

        // Run the model
        UE::NNE::IModelInstanceCPU::ERunSyncStatus RunSyncStatus = ModelHelper->ModelInstance->RunSync(ModelHelper->InputBindings, ModelHelper->OutputBindings);
        if (RunSyncStatus != UE::NNE::IModelInstanceCPU::ERunSyncStatus::Ok)
        {
            UE_LOG(LogTemp, Error, TEXT("Failed to run the model"));
        }

        Mip.BulkData.Unlock();

        AsyncTask(ENamedThreads::GameThread, [WeakThis, ModelHelper]()
        {
            ModelHelper->bIsRunning = false;

            if (WeakThis.IsValid())
            {
                WeakThis->OnEndRunNNE(ModelHelper->OutputData);
            }
        });
    });
}

ブループリント

終わりに

雑な備忘録になりましたが、NNE を使えば ONNX モデルの実行は割と簡単にできるので、あとはそれぞれで必要な ONNX モデルを作っていく感じになりますね。ゲームプレイ向けに ONNX を使って何かというのはなかなかアイデアが出なかったりするのですが、エラーチェッカー的な使い方やアセット生成等には使えそうな場面もあったりするので役立てていけると良さそうです。