本文經(jīng)自動(dòng)駕駛之心公眾號(hào)授權(quán)轉(zhuǎn)載，轉(zhuǎn)載請(qǐng)聯(lián)系出處。

原標(biāo)題：NeRF-XL: Scaling NeRFs with Multiple GPUs

論文鏈接：https://research.nvidia.com/labs/toronto-ai/nerfxl/assets/nerfxl.pdf

項(xiàng)目鏈接：https://research.nvidia.com/labs/toronto-ai/nerfxl/

作者單位：NVIDIA 加州大學(xué)伯克利分校

論文思路：

本文提出了NeRF-XL，這是一種原理性的方法，用于在多個(gè)圖形處理器（GPUs）之間分配神經(jīng)輻射場(chǎng)（NeRFs），從而使得具有任意大容量的NeRF的訓(xùn)練和渲染成為可能。本文首先回顧了現(xiàn)有的多GPU方法，這些方法將大型場(chǎng)景分解成多個(gè)獨(dú)立訓(xùn)練的NeRFs [9, 15, 17]，并確定了這些方法的幾個(gè)基本問題，這些問題在使用額外的計(jì)算資源（GPUs）進(jìn)行訓(xùn)練時(shí)阻礙了重建質(zhì)量的提高。NeRF-XL解決了這些問題，并允許通過簡單地使用更多的硬件來訓(xùn)練和渲染具有任意數(shù)量參數(shù)的NeRFs。本文方法的核心是一個(gè)新穎的分布式訓(xùn)練和渲染公式，這在數(shù)學(xué)上等同于經(jīng)典的單GPU案例，并最小化了GPU之間的通信。通過解鎖具有任意大參數(shù)數(shù)量的NeRFs，本文的方法是第一個(gè)揭示NeRFs多GPU擴(kuò)展規(guī)律(scaling laws)的方法，顯示出隨著參數(shù)數(shù)量的增加而提高的重建質(zhì)量，以及隨著更多GPU的增加而提高的速度。本文在多種數(shù)據(jù)集上展示了NeRF-XL的有效性，包括迄今為止最大的開源數(shù)據(jù)集MatrixCity [5]，它包含了258K張圖像，覆蓋了25平方公里的城市區(qū)域。

論文設(shè)計(jì)：

近期在新視角合成的進(jìn)步極大地提高了我們捕獲神經(jīng)輻射場(chǎng)（NeRFs）的能力，使得這一過程變得更加易于接近。這些進(jìn)步使得我們能夠重建更大的場(chǎng)景和場(chǎng)景內(nèi)更精細(xì)的細(xì)節(jié)。無論是通過增加空間規(guī)模（例如，捕獲數(shù)公里長的城市景觀）還是提高細(xì)節(jié)水平（例如，掃描田野中的草葉），擴(kuò)大捕獲場(chǎng)景的范圍都涉及到將更多的信息量納入NeRF中，以實(shí)現(xiàn)準(zhǔn)確的重建。因此，對(duì)于信息含量高的場(chǎng)景，重建所需的可訓(xùn)練參數(shù)數(shù)量可能會(huì)超過單個(gè)GPU的內(nèi)存容量。

本文提出了NeRF-XL，這是一個(gè)原理性的算法，用于在多個(gè)GPU之間高效分配神經(jīng)輻射場(chǎng)（NeRFs）。本文的方法通過簡單增加硬件資源，使得捕獲高信息含量場(chǎng)景（包括大規(guī)模和高細(xì)節(jié)特征的場(chǎng)景）成為可能。NeRF-XL的核心是在一組不相交的空間區(qū)域之間分配NeRF參數(shù)，并跨GPU聯(lián)合訓(xùn)練它們。不同于傳統(tǒng)的分布式訓(xùn)練流程在后向傳播中同步梯度，本文的方法只需要在前向傳播中同步信息。此外，通過仔細(xì)重寫體渲染方程和分布式設(shè)置中相關(guān)的損失項(xiàng)，本文大幅減少了GPU之間所需的數(shù)據(jù)傳輸。這種新穎的重寫提高了訓(xùn)練和渲染的效率。本文方法的靈活性和可擴(kuò)展性使本文能夠使用多個(gè)GPU高效地優(yōu)化具有任意數(shù)量參數(shù)的NeRFs。

本文的工作與最近采用多GPU算法來建模大規(guī)模場(chǎng)景的方法形成了對(duì)比，這些方法通過訓(xùn)練一組獨(dú)立的NeRFs來實(shí)現(xiàn)[9, 15, 17]。雖然這些方法不需要GPU之間的通信，但每個(gè)NeRF都需要建模整個(gè)空間，包括背景區(qū)域。這導(dǎo)致隨著GPU數(shù)量的增加，模型容量中的冗余度增加。此外，這些方法在渲染時(shí)需要混合NeRFs，這會(huì)降低視覺質(zhì)量并在重疊區(qū)域引入偽影。因此，與NeRF-XL不同的是，這些方法在訓(xùn)練中使用更多的模型參數(shù)（相當(dāng)于更多的GPU）時(shí)，未能實(shí)現(xiàn)視覺質(zhì)量的提升。

本文通過一系列多樣化的捕獲案例來展示本文方法的有效性，包括街道掃描、無人機(jī)飛越和以物體為中心的視頻。這些案例的范圍從小場(chǎng)景（10平方米）到整個(gè)城市（25平方公里）。本文的實(shí)驗(yàn)表明，隨著本文將更多的計(jì)算資源分配給優(yōu)化過程，NeRF-XL始終能夠?qū)崿F(xiàn)改善的視覺質(zhì)量（通過PSNR測(cè)量）和渲染速度。因此，NeRF-XL使得在任何空間規(guī)模和細(xì)節(jié)的場(chǎng)景上訓(xùn)練具有任意大容量的NeRF成為可能。

圖 1：本文基于原理的多GPU分布式訓(xùn)練算法能夠?qū)eRFs擴(kuò)展到任意大的規(guī)模。

圖 2：獨(dú)立訓(xùn)練與多GPU聯(lián)合訓(xùn)練。獨(dú)立地訓(xùn)練多個(gè)NeRFs[9,15,18]要求每個(gè)NeRF既要建模焦點(diǎn)區(qū)域也要建模其周圍環(huán)境，這導(dǎo)致了模型容量的冗余。相比之下，本文的聯(lián)合訓(xùn)練方法使用不重疊的NeRFs，因此沒有任何冗余。

圖 3：獨(dú)立訓(xùn)練需要在新視角合成時(shí)進(jìn)行混合。無論是在2D[9, 15]還是3D[18]中進(jìn)行混合，都會(huì)在渲染中引入模糊。

圖 4：獨(dú)立訓(xùn)練導(dǎo)致不同的相機(jī)優(yōu)化。在NeRF中，相機(jī)優(yōu)化可以通過變換不準(zhǔn)確的相機(jī)本身或所有其他相機(jī)以及底層3D場(chǎng)景來實(shí)現(xiàn)。因此，伴隨相機(jī)優(yōu)化獨(dú)立訓(xùn)練多個(gè)NeRF可能導(dǎo)致相機(jī)校正和場(chǎng)景幾何的不一致性，這給混合渲染帶來了更多困難。

圖 5：3D混合可能造成的視覺偽影。左圖展示了使用2個(gè)GPU訓(xùn)練的MegaNeRF結(jié)果。在0%重疊時(shí)，由于獨(dú)立訓(xùn)練，邊界出現(xiàn)了偽影；而在15%重疊時(shí)，由于3D混合，出現(xiàn)了嚴(yán)重的偽影。右圖闡釋了這種偽影的成因：雖然每個(gè)獨(dú)立訓(xùn)練的NeRF渲染出正確的顏色，但混合后的NeRF并不保證正確的顏色渲染。

圖 6：本文的訓(xùn)練流程。本文的方法聯(lián)合訓(xùn)練所有GPU上的多個(gè)NeRFs，每個(gè)NeRF覆蓋一個(gè)不相交的空間區(qū)域。GPU之間的通信僅發(fā)生在前向傳播中，而不發(fā)生在后向傳播中（如灰色箭頭所示）。(a) 本文可以通過評(píng)估每個(gè)NeRF以獲得樣本顏色和密度，然后將這些值廣播到所有其他GPU以進(jìn)行全局體渲染（見第4.2節(jié)）。(b) 通過重寫體渲染方程，本文可以將數(shù)據(jù)傳輸量大幅減少到每條光線一個(gè)值，從而提高效率（見第4.3節(jié)）。

實(shí)驗(yàn)結(jié)果：

圖 7：定性比較。與先前的工作相比，本文的方法有效地利用多GPU配置，在所有類型的數(shù)據(jù)上提高了性能。

圖 8：定量比較?；讵?dú)立訓(xùn)練的先前工作未能隨著額外GPU的增加而實(shí)現(xiàn)性能提升，而本文的方法隨著訓(xùn)練資源的增加，享受到了渲染質(zhì)量和速度的提升。

圖 9：本文方法的可擴(kuò)展性。更多的GPU允許有更多的可學(xué)習(xí)參數(shù)，這導(dǎo)致了更大的模型容量和更好的質(zhì)量。

圖 10：大規(guī)模捕獲上的更多渲染結(jié)果。本文在更大的捕獲數(shù)據(jù)集上使用更多的GPU測(cè)試了本文方法的魯棒性。請(qǐng)參閱本文的網(wǎng)頁，以獲取這些數(shù)據(jù)的視頻導(dǎo)覽。

圖 11：在University4數(shù)據(jù)集上與PyTorch DDP的比較。PyTorch 分布式數(shù)據(jù)并行（Distributed Data Parallel，DDP）旨在通過跨GPU分布光線來加快渲染速度。相比之下，本文的方法是跨GPU分布參數(shù)，突破了集群中單個(gè)GPU的內(nèi)存限制，并且能夠擴(kuò)大模型容量以獲得更好的質(zhì)量。

圖 12：University4上的同步成本。本文基于分區(qū)的體渲染（見第4.3節(jié)）允許 tile-based 通信，這比原始的基于樣本的通信（見第4.2節(jié)）成本要低得多，因此能夠?qū)崿F(xiàn)更快的渲染。

總結(jié)：

總結(jié)來說，本文重新審視了將大規(guī)模場(chǎng)景分解為獨(dú)立訓(xùn)練的NeRFs（神經(jīng)輻射場(chǎng)）的現(xiàn)有方法，并發(fā)現(xiàn)了阻礙額外計(jì)算資源（GPUs）有效利用的重大問題，這與利用多GPU設(shè)置來提升大規(guī)模NeRF性能的核心目標(biāo)相矛盾。因此，本文引入了NeRF-XL，這是一種原理性的算法，能夠有效地利用多GPU設(shè)置，并通過聯(lián)合訓(xùn)練多個(gè)非重疊的NeRFs來在任何規(guī)模上增強(qiáng)NeRF性能。重要的是，本文的方法不依賴于任何啟發(fā)式規(guī)則，并且在多GPU設(shè)置中遵循NeRF的擴(kuò)展規(guī)律(scaling laws)，適用于各種類型的數(shù)據(jù)。

引用：

@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}
  }