Neural Data Subsystem for unreal engine.

NeuralData.cpp

#include "NeuralData.h"
#include "NNE.h"


UNeuralData::UNeuralData(): SelectedRuntimeType(ERuntimeType::CPU)
{
}

bool UNeuralData::InitModel(const FString& ModelPath, const FString& RuntimeType)
{
	ModelData = LoadObject<UNNEModelData>(GetTransientPackage(), *ModelPath);
	if (RuntimeType == "NNERuntimeORTCuda") {
		SelectedRuntimeType = ERuntimeType::GPU;
		RuntimeGPU = UE::NNE::GetRuntime<INNERuntimeGPU>(RuntimeType);
		auto LoadedModel = RuntimeGPU->CreateModel(ModelData);
		if (LoadedModel && !LoadedModel.IsValid())
		{
			//UE_LOG(LogTemp, Error, TEXT("GPU LoadedModel is invalid"));
			return false;
		}
		ModelInstanceGPU = LoadedModel->CreateModelInstance();
		if (!ModelInstanceGPU.IsValid())
		{
			//UE_LOG(LogTemp, Error, TEXT("GPU ModelInstance is invalid"));
			return false;
		}
	} else if (RuntimeType == "NNERuntimeORTCpu") {
		SelectedRuntimeType = ERuntimeType::CPU;
		RuntimeCPU = UE::NNE::GetRuntime<INNERuntimeCPU>(RuntimeType);
		auto LoadedModel = RuntimeCPU->CreateModel(ModelData);
		if (LoadedModel && !LoadedModel.IsValid())
		{
			//UE_LOG(LogTemp, Error, TEXT("CPU LoadedModel is invalid"));
			return false;
		}
		ModelInstanceCPU = LoadedModel->CreateModelInstance();
		if (!ModelInstanceCPU.IsValid())
		{
			//UE_LOG(LogTemp, Error, TEXT("CPU ModelInstance is invalid"));
			return false;
		}
	} else {
		//UE_LOG(LogTemp, Error, TEXT("Invalid model runtime..."));
		return false;
	}
	return true;
}

bool UNeuralData::IsModelInstanceValid()
{
	switch (SelectedRuntimeType)
	{
	case ERuntimeType::GPU:
		return ModelInstanceGPU.IsValid();
	case ERuntimeType::CPU:
		return ModelInstanceCPU.IsValid();
	default:
		return false;
	}
}
bool UNeuralData::SetShapes(TArray<uint32> InputShape, int32 OutputSize)
{
	ExpectedOutputSize = OutputSize;
	if (!IsModelInstanceValid())
	{
		//UE_LOG(LogTemp, Error, TEXT("Cannot continue, model instance is invalid"));
		return false;
	}
	InputTensorShapes = {UE::NNE::FTensorShape::Make(InputShape)};

	// Set input tensor shapes
	if (SelectedRuntimeType == ERuntimeType::GPU)
	{
		//UE_LOG(LogTemp, Display, TEXT("GPU Model initalizing with: "));
		auto resp = ModelInstanceGPU->SetInputTensorShapes(InputTensorShapes);
		if (resp != 0)
		{
			//UE_LOG(LogTemp, Error, TEXT("Failed to set input tensor shapes"));
			return false;
		}
	} else if (SelectedRuntimeType == ERuntimeType::CPU)
	{
		//UE_LOG(LogTemp, Display, TEXT("CPU Model initalizing with: "));
		auto resp = ModelInstanceCPU->SetInputTensorShapes(InputTensorShapes);
		if (resp != 0)
		{
			//UE_LOG(LogTemp, Error, TEXT("Failed to set input tensor shapes"));
			return false;
		}
	} 
	//UE_LOG(LogTemp, Display, TEXT("Model ready for inference"));
	return true;
}

void UNeuralData::RunClassify(TArray<float> InSeqData, bool inWrapData, uint64_t InStartTimestamp)
{
	FInferenceData Result;
	Result.Timestamp = InStartTimestamp;
	if (bHasCoreIssues)
	{
		return;
	}

	if (!IsModelInstanceValid())
	{
		//UE_LOG(LogTemp, Error, TEXT("Model instance is not valid for inference."));
		return;
	}
	if (bIsRunningInference)
	{
		return;
	}
	bIsRunningInference = true;

	TArray<float> OutputData;
	OutputData.SetNum(ExpectedOutputSize);


	TArray<float> InData;
	if (inWrapData)
	{
		InData.Append(InSeqData);
	} else
	{
		InData = InSeqData;
	}
	
	// Run inference
	int output = 0;
	if (SelectedRuntimeType == ERuntimeType::GPU)
	{
		UE::NNE::FTensorBindingGPU InputBinding;
		InputBinding.Data        = InData.GetData();
		InputBinding.SizeInBytes = InData.Num() * sizeof(float);

		// Set up output tensor binding
		UE::NNE::FTensorBindingGPU OutputBinding;
		OutputBinding.Data        = OutputData.GetData();
		OutputBinding.SizeInBytes = OutputData.Num() * sizeof(int);
		output = ModelInstanceGPU->RunSync({InputBinding}, {OutputBinding});
	} else if (SelectedRuntimeType == ERuntimeType::CPU)
	{
		UE::NNE::FTensorBindingCPU InputBinding;
		InputBinding.Data        = InData.GetData();
		InputBinding.SizeInBytes = InData.Num() * sizeof(float);

		// Set up output tensor binding
		UE::NNE::FTensorBindingCPU OutputBinding;
		OutputBinding.Data        = OutputData.GetData();
		OutputBinding.SizeInBytes = OutputData.Num() * sizeof(int);
		output = ModelInstanceCPU->RunSync({InputBinding}, {OutputBinding});
	}
	if (output != 0)
	{
		//UE_LOG(LogTemp, Error, TEXT("Inference failed."));
		bIsRunningInference = false;
		//Do not broadcast bad results...
		return;
	}
	if (!bHasCoreIssues)
	{
		for (float value : OutputData)
		{
			if (FMath::IsNaN(value))
			{
				bHasCoreIssues = true;
				//UE_LOG(LogTemp, Error, TEXT("NaN value detected in model output, bad input data / model? (InputData.Num() = %d)"), InData.Num());
				//UE_LOG(LogTemp, Warning, TEXT("Output Data is: %f"), *OutputData.GetData());
				// Handle NaN case, e.g., return a default result or take other actions
				bIsRunningInference = false;
				// Do not broadcast NaN results
				return; // Return default FInferenceData
			}
		}
	}

	if (ExpectedOutputSize == 1)
	{
		// Binary classification case
		float probability = Sigmoid(OutputData[0]);
		Result.Cat        = probability > 0.5 ? 1 : 0;
		Result.Confidence = probability;
		Result.bIsValid   = true;
		OnResult.Broadcast(Result);
		return;
	}
	// Multi-class classification case
	TArray<float> Probabilities = Softmax(OutputData);

	// Find the index of the maximum value in the probabilities array
	int MaxIndex   = 0;
	float MaxValue = Probabilities[0];
	for (int i = 1; i < Probabilities.Num(); ++i)
	{
		if (Probabilities[i] > MaxValue)
		{
			MaxValue = Probabilities[i];
			MaxIndex = i;
		}
	}

	// Set the category and confidence
	Result.Cat        = MaxIndex;
	Result.Confidence = MaxValue;
	Result.bIsValid   = true;
	OnResult.Broadcast(Result);
	bIsRunningInference = false;
}

float UNeuralData::Sigmoid(float InLogit)
{
	return 1 / (1 + FMath::Exp(-InLogit));
}

TArray<float> UNeuralData::Softmax(const TArray<float>& InLogits)
{
	TArray<float> expVals;
	float maxLogit = FMath::Max(InLogits); // Find the maximum value for stability
	float sumExpVals = 0.0f;

	for (float logit : InLogits)
	{
		float expVal = FMath::Exp(logit - maxLogit); // Subtract maxLogit for numerical stability
		expVals.Add(expVal);
		sumExpVals += expVal;
	}
	if (sumExpVals == 0) // Check for zero division
	{
		//UE_LOG(LogTemp, Warning, TEXT("Sum of exponential values is zero. Unable to compute softmax."));
		return TArray<float>(); // Return an empty array or handle this case as needed
	}
	for (float& expVal : expVals)
	{
		expVal /= sumExpVals;
	}
	return expVals;
}

NeuralData.h

#pragma once
#include "NNEModelData.h"
#include "UObject/WeakInterfacePtr.h"
#include "TradeMasterX/Core/Inference.h"
#include "NNERuntimeCPU.h"
#include "NNERuntimeGPU.h"
#include "NeuralData.generated.h"

DECLARE_DYNAMIC_MULTICAST_DELEGATE_OneParam(FCurrentResult, FInferenceData, Result);

UCLASS()
class UNeuralData : public UGameInstanceSubsystem
{
	GENERATED_BODY()

public:
	UNeuralData();

	UPROPERTY(BlueprintAssignable, Category = "Data Delegation")
	FCurrentResult OnResult;
	
	UFUNCTION(BlueprintCallable)
	bool InitModel(const FString& ModelPath, const FString& RuntimeType = "NNERuntimeORTCpu");
	bool IsModelInstanceValid();

	bool SetShapes(TArray<uint32> InputShape, int32 OutputSize);

	// Perform inference and return results
	void RunClassify(TArray<float> InData, bool inWrapData, uint64_t InStartTimestamp = 0);
	

private:
	enum class ERuntimeType
	{
		GPU,
		CPU
	};
	ERuntimeType SelectedRuntimeType;
	bool bHasCoreIssues = false;
	UPROPERTY()
	TObjectPtr<UNNEModelData> ModelData;
	TWeakInterfacePtr<INNERuntimeGPU> RuntimeGPU;
	TWeakInterfacePtr<INNERuntimeCPU> RuntimeCPU;
	TUniquePtr<UE::NNE::IModelInstanceGPU> ModelInstanceGPU;
	TUniquePtr<UE::NNE::IModelInstanceCPU> ModelInstanceCPU;
	TArray<float> PreparedData; // Data prepared for inference
	TConstArrayView<UE::NNE::FTensorShape> InputTensorShapes;
	TConstArrayView<UE::NNE::FTensorShape> OutputTensorShapes;
	int32 ExpectedOutputSize = 3;
	bool bIsRunningInference = false;

	TArray<float> Softmax(const TArray<float>& InLogits);
	float Sigmoid(float InLogit);
};