レンダリングを効率化してAIを6倍速にした話

こんにちは、クラスター株式会社メタバース研究所でリサーチエンジニアをしているElizaです。
本記事では先日公開したワールドサムネイル提案システムの基盤となった研究、PanoTree (パノツリー)の特に実装面での工夫についてご紹介したいと思います。
前半はPanoTree技術を簡単に説明し、後半には技術者向けに実装の工夫を詳しく説明しています。

ワールドの素敵な景色を見つけるAI「PanoTree」とは？
- ワールドサムネイル提案システム「WorldPano 」αテスト中です
なぜAIにワールドの魅力を理解させたいのか？
PanoTreeの仕組み
- 全体の流れ
PanoTreeを構成するモジュール
どれくらいの処理量？
レンダリングサーバーの実装
まとめ

ワールドの素敵な景色を見つけるAI「PanoTree」とは？

PanoTreeは撮影スポット自動探索AIの名称で、オープンソースソフトウェアとして公開されています。
PanoTreeはclusterに限らず、Unityで作成された3Dシーンに使用し、探索された撮影スポットをその場で確認することができます。

ワールドサムネイル提案システム「WorldPano 」αテスト中です

現在、PanoTreeを活用したサムネイル提案システム「WorldPano」のαテストを行なっています。
PanoTreeは、クリエイターがCluster Creator Kitで作成したワールドをAIが自動的に探索して、素敵なサムネイルを自動で提案するサービス。もしご自身でワールドをアップロードしているなら、ぜひお試しください。

worldpano.cluster.mu

なぜAIにワールドの魅力を理解させたいのか？

メタバース研究所は「人類の創造力を加速する」というクラスター全体の目標を先導するための組織です。
少し壮大に聞こえるかもしれませんが、VRの新たな可能性を追い求め、クリエイターのみなさんの想像力を引き出すために頑張っている組織くらいに考えてもらえればと思います。

PanoTreeの取り組みは、ワールドの魅力ある景色を探索することを通じて、人を惹きつけるワールドとはとはなにか？をAIに理解させることを目標としています。
ワールドを視覚的に理解しているAIは、クリエイターがワールドを魅力的にするための糸口になります。
今はただ世界を眺めるだけのAIですが、将来的にはクリエイターの想像力を最大限引き出すためのパートナーとなる未来を目指しています。

PanoTreeの仕組み

PanoTreeの特徴は、広大なワールドの探索処理量を大幅に削減し、効率的に探索を行えることです。
この記事では、学術的な言い回しはできるだけ避けて、簡単にその仕組みについて解説します。
もしあなたが研究者なら、arxivも合わせてご覧ください。

全体の流れ

探索する範囲を決める
1. 探索を指示されたPanoTreeはまずワールド全体のコライダーが存在する領域を計算します。
2. これは探索範囲を大まかにコライダーがある場所(=ものがある場所)に絞り込むためです
ワールドを2分割する
1. 例えば、ワールドを左半分と右半分に分割します。
2. 分割を繰り返すことでMinecraftのブロックのように世界が分割されていきます
スクリーンショットを撮影する
1. 分割したそれぞれの領域の中心からカメラを360度ぐるりと回して12枚スクリーンショットを撮影します
景色の魅力を判断する
1. 分割したそれぞれの領域のスクリーンショット12枚の景色の魅力をスコア化し、最も良いスコアをその領域の景色の良さのスコアとします
より魅力的な景色の領域を更に2分割する
1. 分割した後、「3. スクリーンショットの撮影」から処理を繰り返します
一定回数分割し終えたら処理を終了
1. この時点での景色が良いランキング上位のスクリーンショットが、探索結果となります

PanoTreeを構成するモジュール

PanoTreeは大きく分けて3つのモジュールでできています。

レンダリングサーバー: スクリーンショットを撮影するモジュール

このモジュールはUnityで実装されていて、ワールドのスクリーンショットを一度に最大36枚まで効率的に撮影してくれます。

ScoringNet: 景色の魅力を判断するモジュール

PanoTreeの最も重要なパーツです。
このニューラルネットはワールドのスクリーンショット画像を受け取り、それがどれくらい魅力的な景色だったかを0〜100%のスコアで評価します。

HOO: ワールドの探索方法を指示するモジュール

このモジュールは、ワールド領域を分割する際に左右どちらの領域を分割していくべきかを決定します。
ワールドを繰り返し2分割するツリー探索を行うのは、処理量を減らすためです。

どれくらいの処理量？

なぜ処理量を減らす必要があるのか、少し考えてみましょう。
何も工夫しなければどれくらいの処理量になるのでしょうか？

ルービックキューブを思い浮かべてください。
ルービックキューブは縦3、横3、奥行き3ブロックの合計9ブロックです。
9ブロックそれぞれを12視点撮影して評価するのですから、評価すべきスクリーンショットは9 × 12 = 108枚になります。
まだ、なんとかなりそうですね？

では、幅100、奥行き100、高さ100ブロックではどうでしょうか？ブロックは100万個で、1200万枚のスクリーンショットを評価する必要があります。1秒間に1000画像評価したとしても3時間以上かかる計算です。長さが100mを超えるワールドは決して珍しいものではありません。
HOOモジュールが効率的な探索方法を考えてくれるおかげで私達は3時間も待たずに済み、およそ5分で結果を得ることができます。

レンダリングサーバーの実装

さて、ようやく本題のレンダリングサーバーです。ここからは技術者向けの内容になります。

このレンダリングサーバーは前述の通りUnityで実装されています。ここで問題があります。ニューラルネットの性能特性です。ScoringNetはビジョントランスフォーマー(ViT-B/16)をFine-tuningしたモデルです。つまり1枚ずつ画像を与えていたらまったく速度が出ません。複数枚の画像を一度に与え(=バッチサイズを大きくし)なければ処理に何倍も時間がかかります。

従来手法

HOOモジュールで撮影カメラアングル12個を決める
それぞれのカメラアングルについて以下を実行
1. レンダリングサーバーへレンダリングを要求
2. スクリーンショットをレンダリング
3. ScoringNetでスクリーンショットを評価
HOOモジュールが次の探索領域を決定する
1から繰り返す

この方法は仕組みが単純ですが多くの問題があります。

GPUとCPUの同期の頻度が高すぎる
1. レンダリングとScoringNetの計算はそれぞれGPUを使用しますが、この手法ではそれを交互に行っています
  1. GPU上でレンダリング (テクスチャとして)
  2. VRAMからメインメモリに書き戻す
  3. ScoringNetモジュールに転送
  4. メインメモリからVRAMに転送 (pytorchテンソルとして)
2. メインメモリ、VRAM間で安全にデータを転送するにはGPUとCPUが同期を取る必要があります。
3. 1ループごとにGPUとCPUの同期を2回行うため著しくパフォーマンスを低下させます。
GPUでの処理が細切れすぎる
1. 一般にGPUは大量のデータを一度に処理するときに最も性能を発揮するよう設計されています。
2. しかしこの手法では、ScoringNetに1枚ずつ画像が送信されるためGPU本来の性能が発揮されません。
1 Updateに1枚しかスクリーンショットを撮影できない
1. 普通にUnityのスクリーンショット機能を使ってしまうと1 Updateに1スクリーンショットしか撮影できず、60枚/秒です。実際はScoringNetの処理を待機するのでもっと遅くなります。
2. 一方でScoringNetは150枚/秒処理できます。これではレンダリング側の速度が全く足りません。(数字はGPU性能によって変動します)

解決方法はシンプルです。レンダリングもまとめて複数枚行います。ただし、GPUはニューラルネットの計算にも使うのでGPU負荷を極力減らす必要があります。処理を次のように整理します。

改善手法

従来の手法の問題が明らかになってきたところで、それらの問題を解決できるように改善手法の要件を考えてみましょう。

1 Updateにまとめて複数枚のスクリーンショットを撮影する
1. カメラをたくさん配置して 1 Updateでまとめてスクリーンショットを撮影します。
GPUとCPUの同期の頻度を減らす・GPUでの処理をまとめて行う
1. まとめて撮影したスクリーンショットをまとめて処理します。シンプルですね。
2. 12枚の画像を12回VRAM→RAM→VRAMと移動させていたのが1回で済むようになり、同期回数は12分の1に減少します。

実際の手法はこうなります。

HOOモジュールで撮影カメラアングル12個を決める
カメラアングル12個をまとめてレンダリングサーバーに送信しレンダリング
ScoringNetで12枚のスクリーンショットをバッチで評価
HOOモジュールが次の探索領域を決定する
1から繰り返す

改善手法の実装① GPUに優しくないバージョン

実際どうやって1 Updateで12枚のスクリーンショットを撮影するかというと、カメラを12個シーンに配置し、それぞれ結果を読み取るようにします。
しかし素直に実装するとGPUとCPUの同期が12回発生してしまいます。
以下はTextureをメインメモリに読み込むコードです。

public static List<byte[]> TextureReadback(List<Camera> cameras)
{
   var currentRT = RenderTexture.active;
   List<byte[]> results = cameras
       .Select(cam => cam.targetTexture)
       .Select(renderTexture =>
       {
           RenderTexture.active = renderTexture;


           // Texture2D.ReadPixels()によりアクティブなレンダーテクスチャのピクセル情報をテクスチャに格納する
           var texture = new Texture2D(renderTexture.width, renderTexture.height, TextureFormat.RGB24, 0, false);
           texture.ReadPixels(new Rect(0, 0, renderTexture.width, renderTexture.height), 0, 0);
           texture.Apply(); // ここで同期が発生！
           var ret = texture.GetPixelData<byte>(0);
           var binary = new byte[renderTexture.width * renderTexture.height * 3];
           ret.CopyTo(binary);
           return binary;
       }).ToList();


   // アクティブなレンダーテクスチャを元に戻す
   RenderTexture.active = currentRT;
   return results;
}

この例の場合、texture.Apply()した瞬間にGPUとCPUの同期が入ります。
OpenGLでいうとglReadPixelsやglMapBufferのような処理が呼ばれます。(実際には実行OSや描画バックエンド依存です) 同期が増えるとはいえ1 Updateで12カメラ同時に処理できるために従来よりは良いですが、まだ早くなるはずです。
どうにかして同期を減らす方法を考えます。

改善手法の実装② GPUに優しいバージョン

ここでは実際のPanoTreeで採用している実装をご紹介します。

namespace ClusterLab.UseCase.Render
{
   public class TileCameraRenderer: AbstractMultiCameraRenderer
   {
       [SerializeField] RenderTexture tileRenderTexture;


       // 中略


       void InitRenderTexture(int aTextureWidth, int aTextureHeight)
       {
           var targetCamera = Camera.main;
           if (targetCamera == null)
               return;
           // テクスチャの幅と高さ
           rowCol = (int) Math.Ceiling(Math.Sqrt(NumCameras));
           if (textureWidth == aTextureWidth
               && textureHeight == aTextureHeight
               && tileRenderTexture != null
               && tileRenderTexture.width == aTextureWidth
               && tileRenderTexture.height == aTextureHeight)
               return;
           textureWidth = aTextureWidth;
           textureHeight = aTextureHeight;
           renderTextureWidth = rowCol * aTextureWidth;
           renderTextureHeight = rowCol * aTextureHeight;
           // レンダリング結果をタイリングするためのテクスチャ
           tileRenderTexture = new RenderTexture(renderTextureWidth, renderTextureHeight, GraphicsFormat.R8G8B8A8_SRGB, GraphicsFormat.None, 0);
           if (rawImage != null)
               rawImage.texture = tileRenderTexture;


           CommandBuffer.Clear();


           foreach (var y in Enumerable.Range(0, rowCol))
           {
               foreach (var x in Enumerable.Range(0, rowCol))
               {
                   var idx = y * rowCol + x;
                   if (idx > NumCameras - 1)
                       break;


                   // 毎フレームのレンダリング完了後、各カメラのrenderTextureをtileRenderTextureにコピーする命令をcommandBufferに積む
                   CommandBuffer.CopyTexture(
                       AgentCameras[idx].TargetTexture, 0, 0, 0, 0, textureWidth, textureHeight,
                       this.tileRenderTexture, 0, 0, x * textureWidth, textureHeight * (rowCol - 1) - y * textureHeight
                   );
               }
           }
           targetCamera.AddCommandBuffer(CameraEvent.AfterEverything, CommandBuffer);
           isRenderTextureInitialized = true;
       }


       protected override UniTask<List<byte[]>> TextureReadback()
       {
           return tileRenderTexture
               .AsyncGPUReadbackRGB24()
               .ContinueWith(t => new List<byte[]>{t});
       }
   }
}

まずテクスチャを読み取る前に、大きなRenderTextureを1つ用意して12個のテクスチャすべてをRenderTextureにタイリングしていきます。これらの描画はGPUへの事前のコマンド送信によって行われるため、GPUとCPUの同期を回避可能です。
これらのコマンドはメインカメラがアクティブな限り毎フレーム実行されるので無駄が生じますが、GPU内で完結する単なるテクスチャのコピーなので無視できる負荷です。

こうしてタイリングしたRenderTextureの読み取り方にも一工夫加えます。

namespace ClusterLab.Infrastructure.Utils
{
   public static class RenderUtils
   {
       public static UniTask<byte[]> AsyncGPUReadbackRGB24(this RenderTexture renderTexture)
       {
           var source = new UniTaskCompletionSource<byte[]>();
           AsyncGPUReadback.Request(renderTexture, 0, TextureFormat.RGB24, request =>
           {
               var binary = new byte[renderTexture.width * renderTexture.height * 3];
               if (!request.done)
               {
                   source.TrySetException(new Exception("AsyncGPUReadback not done."));
                   Debug.LogError("AsyncGPUReadback not done.");
                   return;
               }
               if (request.hasError)
               {
                   source.TrySetException(new Exception("Unable to read back tileRenderTexture."));
                   Debug.LogError("Unable to read back tileRenderTexture.");
                   return;
               }
               var data = request.GetData<byte>();
               data.CopyTo(binary);
               source.TrySetResult(binary);
           });


           return source.Task;
       }
   }
}

テクスチャのメインメモリへの転送にAsyncGPUReadbackを用いることで、非同期にGPUからテクスチャデータを読み取ります。非同期に読み取ることでドロップフレームを抑止でき、テクスチャデータを待っている間に次のカメラのレンダリングを行うことも可能になりスループットが改善します。

どれくらい早くなるの？

PanoTreeの処理全体で、およそ6倍の高速化を達成しました。(Windows/DirectX/Geforce RTX4060の場合)

まとめ

Unityを使った機械学習アプリケーションで、レンダリング画像を効率的にニューラルネットに転送する方法を紹介した
Unityを使った機械学習アプリケーションではCPUとGPUの同期を減らすこと、機械学習部分とUnityでGPU負荷をうまく分け合うように考慮すると高速化できることと、具体的な方法例を紹介
この手法は機械学習アプリケーションに限らず活用できる