Unity Barracudaを用いたIn-Game Style Transferの導入メモとパフォーマンス調査

はじめに

本記事ではUnity Brracudaを用いたIn-Game Style Transferの導入とパフォーマンス調査を行う．基本的に以下のチュートリアルPart3の内容を追っていくが，画像付きで手順が記されているので本記事では手順をすべて解説せず，重要だと思ったところをメモする．また，ソースコードを一部改変している部分もあるため，元記事を一読することを推奨する．
https://www.intel.com/content/www/us/en/developer/articles/training/in-game-style-transfer-leveraging-unity-part-3.html

なお，このチュートリアルを一から進めたい場合はこちらから進めるとよい．

最終的な出力

先にどのような結果が得られるかを以下に示す．

• 左側：入力画面
• 右側：Style Transferの出力画面
  
  Style Tranferの実行結果(gif)

Style Tranferの実行結果(png)

実行環境

• Unity
  • Version: 2019.4.20f1
  • Project: Kinematica_Demo_0.8.0-preview
    • Link: https://christianjmills.com/posts/end-to-end-in-game-style-transfer-tutorial/part-1/#download-kinematica-demo
  • Barracuda: 1.0.4
• PC
  • OS: Windows11
  • CPU: i9-10850K CPU @ 3.60GHz 3.50 GHz
  • GPU: NVIDIA GeForce RTX 3090
  • メモリ: 64.0 GB
• Style Transferモデル
  • Input size: [1, 3, 960, 540]
  • Output size: [1, 3, 960, 540]

想定読者

• Unityを触ったことがある
• Compute Shaderを触ったことがある
• 機械学習の基本的な知識がある

導入メモ

1. Compute Shaderの作成

前処理編

早速，Style TransferモデルにUnityのスクリーン画像を入力したいが，そのまま入力すると色の取りうる値で問題が起きる．Unityでは[0, 1]であるのに対して，Style Transferモデルでは[0, 255]であるため^[1]，ここを調整する必要がある．この調整（前処理）をGPUで行いたいので，Compute Shaderを利用する．
以下のような関数を作り，色のスケーリングを行う．

// Each #kernel tells which function to compile; you can have many kernels
#pragma kernel ProcessInput

// Create a RenderTexture with enableRandomWrite flag and set it
// with cs.SetTexture
RWTexture2D<half4> Result;
Texture2D<half4> InputImage;

[numthreads(8,8,1)]
void ProcessInput (uint3 id : SV_DispatchThreadID)
{
    Result[id.xy] = half4((InputImage[id.xy].x * 255.0h),
        (InputImage[id.xy].y * 255.0h),
        (InputImage[id.xy].z * 255.0h), 1.0h);
}

後処理編

前処理と同様にUnityとStyle Transferモデルの色の範囲が異なるので調整を行うが，時々Style Transferモデルの出力が[0, 255]の範囲を超える場合があるらしい．そこはclamp()を用いることで0から255の範囲に収まるようになる．
なお，ProcessOutputのif (id.x < 480)は画面左側の描画処理を表しており，elseは画面右側の描画処理を表している．

// Each #kernel tells which function to compile; you can have many kernels
#pragma kernel ProcessOutput

// Create a RenderTexture with enableRandomWrite flag and set it
// with cs.SetTexture
RWTexture2D<half4> Result;
Texture2D<half4> InputImage;
Texture2D<half4> OutputImage;

[numthreads(8, 8, 1)]
void ProcessOutput(uint3 id : SV_DispatchThreadID)
{
    if (id.x < 480)
    {
        Result[id.xy] = half4(InputImage[id.xy].xyz, 1.0h);
    }
    else
    {
        Result[id.xy] = half4((clamp(OutputImage[id.xy].x, 0.0f, 255.0f) / 255.0f),
            (clamp(OutputImage[id.xy].y, 0.0f, 255.0f) / 255.0f),
            (clamp(OutputImage[id.xy].z, 0.0f, 255.0f) / 255.0f), 1.0h);
    }
}

2. StyleTransferスクリプトの作成

全てのC#側のコードを説明しないが，上記のCompute Shaderの後処理については自作部分も含まれているため，そこに対応するC#側の後処理関数の説明を行う．
C#側の後処理関数では以下の流れで処理を行う．

1. Compute Shaderのインデックスを取得
2. Shaderの処理結果を書き込む一時的なレンダーテクスチャを作成
3. Compute Shaderの実行
4. 処理結果の画像を入力に使用したレンダーテクスチャにコピー
5. 一時的なレンダーテクスチャをリリース

チュートリアルではProcessImage()のみだったが，自作では前処理と後処理を分けて実装した．以下の後処理関数では入力画像と出力画像を受け取り，それらをShaderに渡すようしている．

private void PostProcessing(RenderTexture outputTexture, RenderTexture inputTexture, string functionName)
{
	int numthreads = 8;
	int kernelHandle = styleTransferShader.FindKernel(functionName);

	RenderTexture result = RenderTexture.GetTemporary(outputTexture.width, outputTexture.height, 24, RenderTextureFormat.ARGBHalf);
	result.enableRandomWrite = true;
	result.Create();

	styleTransferShader.SetTexture(kernelHandle, "Result", result);
	styleTransferShader.SetTexture(kernelHandle, "OutputImage", outputTexture);
	styleTransferShader.SetTexture(kernelHandle, "InputImage", inputTexture);

	styleTransferShader.Dispatch(kernelHandle, result.width / numthreads, result.height / numthreads, 1);
	Graphics.Blit(result, outputTexture);
	RenderTexture.ReleaseTemporary(result);
}

前処理とStylizeImage関数も含めて書き出すと以下のようになる．

private void PreProcessing(RenderTexture inputTexture, string functionName)
{
	int numthreads = 8;
	int kernelHandle = styleTransferShader.FindKernel(functionName);

	RenderTexture result = RenderTexture.GetTemporary(inputTexture.width, inputTexture.height, 24, RenderTextureFormat.ARGBHalf);
	result.enableRandomWrite = true;
	result.Create();

	styleTransferShader.SetTexture(kernelHandle, "Result", result);
	styleTransferShader.SetTexture(kernelHandle, "InputImage", inputTexture);

	styleTransferShader.Dispatch(kernelHandle, result.width / numthreads, result.height / numthreads, 1);
	Graphics.Blit(result, inputTexture);
	RenderTexture.ReleaseTemporary(result);
}

private void PostProcessing(RenderTexture outputTexture, RenderTexture inputTexture, string functionName)
{
	int numthreads = 8;
	int kernelHandle = styleTransferShader.FindKernel(functionName);

	RenderTexture result = RenderTexture.GetTemporary(outputTexture.width, outputTexture.height, 24, RenderTextureFormat.ARGBHalf);
	result.enableRandomWrite = true;
	result.Create();

	styleTransferShader.SetTexture(kernelHandle, "Result", result);
	styleTransferShader.SetTexture(kernelHandle, "OutputImage", outputTexture);
	styleTransferShader.SetTexture(kernelHandle, "InputImage", inputTexture);

	styleTransferShader.Dispatch(kernelHandle, result.width / numthreads, result.height / numthreads, 1);
	Graphics.Blit(result, outputTexture);
	RenderTexture.ReleaseTemporary(result);
}

private void StylizeImage(RenderTexture src)
{
	RenderTexture rTex;
	RenderTexture rTexSrc;

	if (src.height > targetHeight && targetHeight >= 8)
	{
	    //float scale = src.height / targetHeight;
	    //int targetWidth = (int)(src.width / scale);

	    //targetHeight -= (targetHeight % 8);
	    //targetWidth -= (targetWidth & 8);

	    int targetWidth = 960;

	    rTex = RenderTexture.GetTemporary(targetWidth, targetHeight, 24, src.format);
	    rTexSrc = RenderTexture.GetTemporary(targetWidth, targetHeight, 24, src.format);
	}
	else
	{
	    rTex = RenderTexture.GetTemporary(src.width, src.height, 24, src.format);
	    rTexSrc = RenderTexture.GetTemporary(src.width, src.height, 24, src.format);
	}

	Graphics.Blit(src, rTex);
	PreProcessing(rTex, "ProcessInput");
	Tensor input = new Tensor(rTex, channels: 3);

	m_engine.Execute(input);
	Tensor prediction = m_engine.PeekOutput();
	input.Dispose();

	RenderTexture.active = null;
	prediction.ToRenderTexture(rTex);
	prediction.Dispose();

	Graphics.Blit(src, rTexSrc);
	PostProcessing(rTex, rTexSrc, "ProcessOutput");
	Graphics.Blit(rTex, src);
	RenderTexture.ReleaseTemporary(rTex);
	RenderTexture.ReleaseTemporary(rTexSrc);
}

パフォーマンス調査

UnityのProfilerを使ってCPU側とGPU側のBarracudaの実行処理について調査を行う．

CPU

GPU

所感

推論自体はGPUで行っており，FPSも95～110くらいなのでそれほど重いわけではない．ただ，Style Transferモデルのアーキテクチャの見直しやONNXの最適化などでパフォーマンス改善が期待できると思う．また，ちらつきの抑制に取り組むことも面白いと思う．

参考

https://github.com/cj-mills/End-to-End-In-Game-Style-Transfer-Tutorial-Intel
https://light11.hatenadiary.com/entry/2019/09/19/000006
https://christianjmills.com/posts/end-to-end-in-game-style-transfer-tutorial/part-3/index.html

脚注

1. https://github.com/onnx/models/tree/main/vision/style_transfer/fast_neural_style#input-to-model