計圖:五秒訓(xùn)好NeRF!已開源
本文經(jīng)AI新媒體量子位(公眾號ID:QbitAI)授權(quán)轉(zhuǎn)載,轉(zhuǎn)載請聯(lián)系出處。
計圖(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并不僅僅是因為哈希編碼方法,更因為Nvidia針對硬件做的極致優(yōu)化。
△ 圖4:Instant-NGP 算法流程
Instant-NGP能5秒訓(xùn)練NeRF的效果給學(xué)術(shù)界和產(chǎn)業(yè)界都帶來了很多可能,但I(xiàn)nstant-NGP的實現(xiàn)仍存在一些問題:其源碼完全基于Cuda編寫,對習(xí)慣Python的用戶而言使用門檻較高。
Instant-NGP沒有成熟的深度學(xué)習(xí)框架支持,無法適配各種常用NeRF變種模型;Nvidia為了追求極致的效率優(yōu)化,其源碼不同函數(shù)間耦合嚴(yán)重,難以修改。
基于JNeRF實現(xiàn)的Instant-NGP具有以下優(yōu)勢:
- 準(zhǔn)確復(fù)現(xiàn)了論文的速度和點數(shù)(見表1),其他深度學(xué)習(xí)框架復(fù)現(xiàn)的Instant-NGP在速度和點數(shù)上都距離原論文有一定差距,Jittor是全球唯一成功復(fù)現(xiàn)了Instant-NGP的框架。
△ 表1:與Instant NGP原文的對比
- JNeRF速度十分高效,其訓(xùn)練速度可以達(dá)到大約133 iter/s。我們注意到,隨著NeRF訓(xùn)練速度的提升,框架的運行速度成為了限制NeRF速度進(jìn)一步提升的瓶頸,近期有的工作(如Plenoxel)通過大量修改Pytorch源碼才實現(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都無法實現(xiàn)的操作,例如在訓(xùn)練時動態(tài)修改Cuda代碼。
- JNeRF進(jìn)行了較好的模塊化,使得科研人員能夠輕松地調(diào)用、替換其中的模塊。
- Jittor基于元算子融合的優(yōu)勢,有著豐富的硬件生態(tài)支持,因此基于Jittor的JNeRF庫能夠輕松地在各種國內(nèi)外硬件以及移動端上運行。
JNeRF代碼演示
JNeRF通過配置文件可以輕松修改模型設(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