計圖：5秒訓(xùn)好NeRF！已開源

十三 2022-06-05 13:52:42 來源：量子位

金磊 轉(zhuǎn)載整理自 圖形學(xué)與幾何計算

量子位 | 公眾號 QbitAI

計圖（Jittor）框架的NeRF模型庫JNeRF正式開源了！

通過JNeRF可以5秒訓(xùn)練好NeRF模型（見圖1）。

Jittor成為首個支持Instant-NGP的深度學(xué)習(xí)框架。

△ 圖1：5秒訓(xùn)練好NeRF

研究背景

NeRF在2020[1]被提出，它驚艷的視角生成效果獲得了學(xué)術(shù)界的廣泛關(guān)注。

不同于傳統(tǒng)的顯式表達(dá)，NeRF通過將場景表示為隱式的神經(jīng)輻射場，渲染時通過神經(jīng)網(wǎng)絡(luò)查詢位置上的場景信息生成新視角圖像。

△ 圖2：用Jittor框架訓(xùn)練的NeRF的例子

NeRF的出現(xiàn)很可能帶來一場革命，但它仍存在一些問題。

不同方法針對NeRF各部分（如：采樣、位置編碼、網(wǎng)絡(luò)結(jié)構(gòu)）的優(yōu)化都會對NeRF的性能帶來提升，這一定程度上影響了比較的公平性。

并且隨著NeRF效率的提升，深度學(xué)習(xí)框架逐漸成為NeRF效率提升的瓶頸，許多工作修改、添加cuda代碼以提升效率，這進(jìn)一步帶來了混亂。

因此NeRF需要一個統(tǒng)一、高效、并與框架高度適配的模型庫，建立統(tǒng)一流程，提升科研效率。

JNeRF架構(gòu)

JNeRF模型庫分析了現(xiàn)有的NeRF方法，將NeRF的主要訓(xùn)練流程總結(jié)為圖3所示的7個模塊，不同模塊之間相互解耦，使其能夠方便地調(diào)用和替換。

由于時間有限，目前JNeRF只支持了其中很少的工作，后續(xù)JNeRF會不斷支持更多有代表性的NeRF工作，也十分歡迎大家在GitHub為JNeRF貢獻(xiàn)代碼。

△ 圖3：NeRF的訓(xùn)練流程

JNeRF支持Instant-NGP

Nvidia的工作Instant-NGP[2]在今年年初發(fā)布后，受到了學(xué)術(shù)界和工業(yè)界的廣泛關(guān)注，github開源倉庫至今已獲得7k+個星。

該工作通過Hash編碼和定制化的優(yōu)化，使得NeRF能在5秒訓(xùn)練出一個效果優(yōu)質(zhì)的結(jié)果。

經(jīng)過我們的分析，Instant-NGP能5秒訓(xùn)練出NeRF并不僅僅是因?yàn)楣＞幋a方法，更因?yàn)镹vidia針對硬件做的極致優(yōu)化。

△ 圖4：Instant-NGP 算法流程

Instant-NGP能5秒訓(xùn)練NeRF的效果給學(xué)術(shù)界和產(chǎn)業(yè)界都帶來了很多可能，但I(xiàn)nstant-NGP的實(shí)現(xiàn)仍存在一些問題：其源碼完全基于Cuda編寫，對習(xí)慣Python的用戶而言使用門檻較高。

Instant-NGP沒有成熟的深度學(xué)習(xí)框架支持，無法適配各種常用NeRF變種模型；Nvidia為了追求極致的效率優(yōu)化，其源碼不同函數(shù)間耦合嚴(yán)重，難以修改。

基于JNeRF實(shí)現(xiàn)的Instant-NGP具有以下優(yōu)勢：

• 準(zhǔn)確復(fù)現(xiàn)了論文的速度和點(diǎn)數(shù)（見表1），其他深度學(xué)習(xí)框架復(fù)現(xiàn)的Instant-NGP在速度和點(diǎn)數(shù)上都距離原論文有一定差距，Jittor是全球唯一成功復(fù)現(xiàn)了Instant-NGP的框架。

△ 表1：與Instant NGP原文的對比

• JNeRF速度十分高效，其訓(xùn)練速度可以達(dá)到大約133 iter/s。我們注意到，隨著NeRF訓(xùn)練速度的提升，框架的運(yùn)行速度成為了限制NeRF速度進(jìn)一步提升的瓶頸，近期有的工作（如Plenoxel）通過大量修改Pytorch源碼才實(shí)現(xiàn)快速訓(xùn)練，而Instant-NGP則更是直接使用Cuda來達(dá)到速度的極致追求。得益于Jittor的速度優(yōu)勢，JNeRF可以直接輕松地使用Jittor框架完成Instant-NGP的復(fù)現(xiàn)。
• Jittor是一個即時編譯的框架，通過即時編譯的特性，能夠?qū)崿F(xiàn)很多原版Instant-NGP都無法實(shí)現(xiàn)的操作，例如在訓(xùn)練時動態(tài)修改Cuda代碼。
• JNeRF進(jìn)行了較好的模塊化，使得科研人員能夠輕松地調(diào)用、替換其中的模塊。
• Jittor基于元算子融合的優(yōu)勢，有著豐富的硬件生態(tài)支持，因此基于Jittor的JNeRF庫能夠輕松地在各種國內(nèi)外硬件以及移動端上運(yùn)行。

JNeRF代碼演示

JNeRF通過配置文件可以輕松修改模型設(shè)置，下面的代碼展示了用戶可以通過簡單地修改配置文件實(shí)現(xiàn)模塊的替換。

目前，JNeRF倉庫已在github和中國計算機(jī)學(xué)會的GitLink上開源。JNeRF仍有待不斷完善，支持的模型仍十分有限。

歡迎大家通過該倉庫進(jìn)行NeRF的學(xué)習(xí)和研究，并和我們一起完善該倉庫，共同推動NeRF的發(fā)展。

GitHub鏈接：

https://github.com/Jittor/JNeRF

GitLink鏈接：

https://www.gitlink.org.cn/jittor/jnerf

參考文獻(xiàn)：

[1]Mildenhall B, Srinivasan P P, Tancik M, et al. Nerf: Representing scenes as neural radiance fields for view synthesis[C]//European conference on computer vision. Springer, Cham, 2020: 405-421.
[2]MüllerT, Evans A, Schied C, et al., Instant Neural Graphics Primitives with a Multiresolution Hash Encoding, arXiv: 2201.05989, 2022.