計圖:五秒訓好NeRF!已開源
本文經AI新媒體量子位(公眾號ID:QbitAI)授權轉載,轉載請聯系出處。
計圖(Jittor)框架的NeRF模型庫JNeRF正式開源了!
通過JNeRF可以5秒訓練好NeRF模型(見圖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開源倉庫至今已獲得7k+個星。
該工作通過Hash編碼和定制化的優化,使得NeRF能在5秒訓練出一個效果優質的結果。
經過我們的分析,Instant-NGP能5秒訓練出NeRF并不僅僅是因為哈希編碼方法,更因為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是一個即時編譯的框架,通過即時編譯的特性,能夠實現很多原版Instant-NGP都無法實現的操作,例如在訓練時動態修改Cuda代碼。
- JNeRF進行了較好的模塊化,使得科研人員能夠輕松地調用、替換其中的模塊。
- Jittor基于元算子融合的優勢,有著豐富的硬件生態支持,因此基于Jittor的JNeRF庫能夠輕松地在各種國內外硬件以及移動端上運行。
JNeRF代碼演示
JNeRF通過配置文件可以輕松修改模型設置,下面的代碼展示了用戶可以通過簡單地修改配置文件實現模塊的替換。
目前,JNeRF倉庫已在github和中國計算機學會的GitLink上開源。JNeRF仍有待不斷完善,支持的模型仍十分有限。
歡迎大家通過該倉庫進行NeRF的學習和研究,并和我們一起完善該倉庫,共同推動NeRF的發展。
GitHub鏈接:
https://github.com/Jittor/JNeRF
GitLink鏈接:
https://www.gitlink.org.cn/jittor/jnerf