Jitu: 5秒でNeRFをトレーニングしましょう!オープンソース

この記事はAI新メディアQuantum Bit（公開アカウントID：QbitAI）より許可を得て転載しています。転載の際は出典元にご連絡ください。

Jittorフレームワークの NeRF モデルライブラリJNeRF が正式にオープンソースになりました。

NeRFモデルはJNeRFを使用して5秒でトレーニングできます（図1を参照）。

Jittor は、Instant-NGP をサポートする最初のディープラーニングフレームワークになります。

△ 図1: 5秒でNeRFトレーニング

背景

NeRFは2020年に提案され[1]、その驚くべき視点生成効果は学術界で広く注目を集めています。

従来の明示的な表現とは異なり、NeRF はシーンを暗黙的なニューラル放射フィールドとして表現し、レンダリング中にニューラルネットワークを通じてその場所のシーン情報を照会することで、新しい遠近法の画像を生成します。

図2: JittorフレームワークでトレーニングされたNeRFの例

NeRF の出現は革命をもたらす可能性は高いが、まだいくつかの問題が残っている。

NeRF のさまざまな部分(サンプリング、位置エンコーディング、ネットワーク構造など) を最適化するさまざまな方法により、NeRF のパフォーマンスが向上し、比較の公平性にある程度影響します。

そして、NeRF の効率が向上するにつれて、ディープラーニングフレームワークが徐々に NeRF の効率向上のボトルネックになってきました。多くの作業者が効率向上のために cuda コードを修正、追加し、さらに混乱を招いています。

したがって、NeRF では、統一されたプロセスを確立し、科学研究の効率を向上させるために、統一され、効率的で、互換性の高いモデルライブラリが必要です。

JNeRFアーキテクチャ

JNeRF モデルライブラリは、既存の NeRF メソッドを分析し、NeRF の主なトレーニングプロセスを図 3 に示すように7 つのモジュールにまとめます。異なるモジュールは互いに分離されているため、呼び出しや置き換えが容易です。

時間の制約により、JNeRF は現在、これらの作品のうちのごく一部しかサポートしていません。今後、JNeRF は、より多くの代表的な NeRF 作品のサポートを継続していきます。また、GitHub で JNeRF にコードを投稿していただくことも歓迎します。

△ 図3: NeRFトレーニングプロセス

JNeRFはInstant-NGPをサポート

今年初めにNvidiaのInstant-NGP[2]がリリースされて以来、学界や産業界から幅広い注目を集めており、githubのオープンソースリポジトリは現在までに7,000以上のスターを獲得しています。

この研究では、ハッシュコーディングとカスタマイズされた最適化を使用して、NeRF が5 秒で高品質の結果をトレーニングできるようにしました。

分析の結果、Instant-NGP が NeRF を 5 秒でトレーニングできるのは、ハッシュコーディング方式だけでなく、Nvidia のハードウェアの極端な最適化によるものであることがわかりました。

図4: Instant-NGPアルゴリズムのフロー

Instant-NGP が 5 秒で NeRF をトレーニングできるという事実は、学界と産業界の両方に多くの可能性をもたらしましたが、Instant-NGP の実装にはまだいくつかの問題があります。そのソースコードは完全に Cuda ベースで記述されており、Python に慣れているユーザーにとっては敷居が高いのです。

Instant-NGP は成熟したディープラーニングフレームワークのサポートがなく、一般的に使用されているさまざまな NeRF バリアントモデルに適応できません。極限の効率最適化を追求するために、Nvidia のソースコードはさまざまな機能間の結合が厳しく、変更が困難になっています。

JNeRF に基づく Instant-NGP には、次のような利点があります。

論文の速度とポイントを正確に再現しました（表1参照）。他のディープラーニングフレームワークで再現されたInstant-NGPは、元の論文とは速度とポイントに一定のギャップがあります。Jittorは、Instant-NGPの再現に成功した世界で唯一のフレームワークです。

△ 表1: オリジナルのInstant NGPテキストとの比較

JNeRF は非常に効率的で、約 133 iter/s の速度でトレーニングできます。 NeRF のトレーニング速度が上がるにつれて、フレームワークの実行速度がボトルネックとなり、NeRF 速度のさらなる向上が制限されることがわかりました。最近のいくつかの研究(Plenoxel など)では、Pytorch ソースコードを大幅に変更することで高速トレーニングを実現しましたが、Instant-NGP では Cuda を直接使用して究極の速度の追求を実現しました。 Jittor の速度の利点により、JNeRF は Jittor フレームワークを直接かつ簡単に使用して Instant-NGP を再現できます。
Jittor はジャストインタイムコンパイルフレームワークです。ジャストインタイムコンパイル機能により、トレーニング中に Cuda コードを動的に変更するなど、元の Instant-NGP では実現できなかった多くの操作を実装できます。
JNeRF は適切にモジュール化されており、研究者はモジュールを簡単に呼び出して交換できます。
Jittor は、メタ演算子の融合と豊富なハードウェアエコロジカルサポートの利点を備えているため、Jittor ベースの JNeRF ライブラリは、国内外のさまざまなハードウェアやモバイル端末で簡単に実行できます。