史上最大重建25km2!NeRF-XL:真正有效利用多卡聯合訓練!
本文經自動駕駛之心公眾號授權轉載,轉載請聯系出處。
原標題:NeRF-XL: Scaling NeRFs with Multiple GPUs
論文鏈接:https://research.nvidia.com/labs/toronto-ai/nerfxl/assets/nerfxl.pdf
項目鏈接:https://research.nvidia.com/labs/toronto-ai/nerfxl/
作者單位:NVIDIA 加州大學伯克利分校
論文思路:
本文提出了NeRF-XL,這是一種原理性的方法,用于在多個圖形處理器(GPUs)之間分配神經輻射場(NeRFs),從而使得具有任意大容量的NeRF的訓練和渲染成為可能。本文首先回顧了現有的多GPU方法,這些方法將大型場景分解成多個獨立訓練的NeRFs [9, 15, 17],并確定了這些方法的幾個基本問題,這些問題在使用額外的計算資源(GPUs)進行訓練時阻礙了重建質量的提高。NeRF-XL解決了這些問題,并允許通過簡單地使用更多的硬件來訓練和渲染具有任意數量參數的NeRFs。本文方法的核心是一個新穎的分布式訓練和渲染公式,這在數學上等同于經典的單GPU案例,并最小化了GPU之間的通信。通過解鎖具有任意大參數數量的NeRFs,本文的方法是第一個揭示NeRFs多GPU擴展規律(scaling laws)的方法,顯示出隨著參數數量的增加而提高的重建質量,以及隨著更多GPU的增加而提高的速度。本文在多種數據集上展示了NeRF-XL的有效性,包括迄今為止最大的開源數據集MatrixCity [5],它包含了258K張圖像,覆蓋了25平方公里的城市區域。
論文設計:
近期在新視角合成的進步極大地提高了我們捕獲神經輻射場(NeRFs)的能力,使得這一過程變得更加易于接近。這些進步使得我們能夠重建更大的場景和場景內更精細的細節。無論是通過增加空間規模(例如,捕獲數公里長的城市景觀)還是提高細節水平(例如,掃描田野中的草葉),擴大捕獲場景的范圍都涉及到將更多的信息量納入NeRF中,以實現準確的重建。因此,對于信息含量高的場景,重建所需的可訓練參數數量可能會超過單個GPU的內存容量。
本文提出了NeRF-XL,這是一個原理性的算法,用于在多個GPU之間高效分配神經輻射場(NeRFs)。本文的方法通過簡單增加硬件資源,使得捕獲高信息含量場景(包括大規模和高細節特征的場景)成為可能。NeRF-XL的核心是在一組不相交的空間區域之間分配NeRF參數,并跨GPU聯合訓練它們。不同于傳統的分布式訓練流程在后向傳播中同步梯度,本文的方法只需要在前向傳播中同步信息。此外,通過仔細重寫體渲染方程和分布式設置中相關的損失項,本文大幅減少了GPU之間所需的數據傳輸。這種新穎的重寫提高了訓練和渲染的效率。本文方法的靈活性和可擴展性使本文能夠使用多個GPU高效地優化具有任意數量參數的NeRFs。
本文的工作與最近采用多GPU算法來建模大規模場景的方法形成了對比,這些方法通過訓練一組獨立的NeRFs來實現[9, 15, 17]。雖然這些方法不需要GPU之間的通信,但每個NeRF都需要建模整個空間,包括背景區域。這導致隨著GPU數量的增加,模型容量中的冗余度增加。此外,這些方法在渲染時需要混合NeRFs,這會降低視覺質量并在重疊區域引入偽影。因此,與NeRF-XL不同的是,這些方法在訓練中使用更多的模型參數(相當于更多的GPU)時,未能實現視覺質量的提升。
本文通過一系列多樣化的捕獲案例來展示本文方法的有效性,包括街道掃描、無人機飛越和以物體為中心的視頻。這些案例的范圍從小場景(10平方米)到整個城市(25平方公里)。本文的實驗表明,隨著本文將更多的計算資源分配給優化過程,NeRF-XL始終能夠實現改善的視覺質量(通過PSNR測量)和渲染速度。因此,NeRF-XL使得在任何空間規模和細節的場景上訓練具有任意大容量的NeRF成為可能。
圖 1:本文基于原理的多GPU分布式訓練算法能夠將NeRFs擴展到任意大的規模。
圖 2:獨立訓練與多GPU聯合訓練。獨立地訓練多個NeRFs[9,15,18]要求每個NeRF既要建模焦點區域也要建模其周圍環境,這導致了模型容量的冗余。相比之下,本文的聯合訓練方法使用不重疊的NeRFs,因此沒有任何冗余。
圖 3:獨立訓練需要在新視角合成時進行混合。無論是在2D[9, 15]還是3D[18]中進行混合,都會在渲染中引入模糊。
圖 4:獨立訓練導致不同的相機優化。在NeRF中,相機優化可以通過變換不準確的相機本身或所有其他相機以及底層3D場景來實現。因此,伴隨相機優化獨立訓練多個NeRF可能導致相機校正和場景幾何的不一致性,這給混合渲染帶來了更多困難。
圖 5:3D混合可能造成的視覺偽影。左圖展示了使用2個GPU訓練的MegaNeRF結果。在0%重疊時,由于獨立訓練,邊界出現了偽影;而在15%重疊時,由于3D混合,出現了嚴重的偽影。右圖闡釋了這種偽影的成因:雖然每個獨立訓練的NeRF渲染出正確的顏色,但混合后的NeRF并不保證正確的顏色渲染。
圖 6:本文的訓練流程。本文的方法聯合訓練所有GPU上的多個NeRFs,每個NeRF覆蓋一個不相交的空間區域。GPU之間的通信僅發生在前向傳播中,而不發生在后向傳播中(如灰色箭頭所示)。(a) 本文可以通過評估每個NeRF以獲得樣本顏色和密度,然后將這些值廣播到所有其他GPU以進行全局體渲染(見第4.2節)。(b) 通過重寫體渲染方程,本文可以將數據傳輸量大幅減少到每條光線一個值,從而提高效率(見第4.3節)。
實驗結果:
圖 7:定性比較。與先前的工作相比,本文的方法有效地利用多GPU配置,在所有類型的數據上提高了性能。
圖 8:定量比較?;讵毩⒂柧毜南惹肮ぷ魑茨茈S著額外GPU的增加而實現性能提升,而本文的方法隨著訓練資源的增加,享受到了渲染質量和速度的提升。
圖 9:本文方法的可擴展性。更多的GPU允許有更多的可學習參數,這導致了更大的模型容量和更好的質量。
圖 10:大規模捕獲上的更多渲染結果。本文在更大的捕獲數據集上使用更多的GPU測試了本文方法的魯棒性。請參閱本文的網頁,以獲取這些數據的視頻導覽。
圖 11:在University4數據集上與PyTorch DDP的比較。PyTorch 分布式數據并行(Distributed Data Parallel,DDP)旨在通過跨GPU分布光線來加快渲染速度。相比之下,本文的方法是跨GPU分布參數,突破了集群中單個GPU的內存限制,并且能夠擴大模型容量以獲得更好的質量。
圖 12:University4上的同步成本。本文基于分區的體渲染(見第4.3節)允許 tile-based 通信,這比原始的基于樣本的通信(見第4.2節)成本要低得多,因此能夠實現更快的渲染。
總結:
總結來說,本文重新審視了將大規模場景分解為獨立訓練的NeRFs(神經輻射場)的現有方法,并發現了阻礙額外計算資源(GPUs)有效利用的重大問題,這與利用多GPU設置來提升大規模NeRF性能的核心目標相矛盾。因此,本文引入了NeRF-XL,這是一種原理性的算法,能夠有效地利用多GPU設置,并通過聯合訓練多個非重疊的NeRFs來在任何規模上增強NeRF性能。重要的是,本文的方法不依賴于任何啟發式規則,并且在多GPU設置中遵循NeRF的擴展規律(scaling laws),適用于各種類型的數據。
引用:
@misc{li2024nerfxl,
title={NeRF-XL: Scaling NeRFs with Multiple GPUs},
author={Ruilong Li and Sanja Fidler and Angjoo Kanazawa and Francis Williams},
year={2024},
eprint={2404.16221},
archivePrefix={arXiv},
primaryClass={cs.CV}
}