首頁>編程相關>其他源碼
下載

微型 CUDA 神經網絡

這是一個小型的、獨立的框架，用於訓練和查詢神經網路。最值得注意的是，它包含閃電般快速的「完全融合」多層感知器（技術論文）、多功能多重解析度雜湊編碼（技術論文），以及對各種其他輸入編碼、損失和優化器的支持。

表現

影像完全融合的網路與帶有 XLA 的 TensorFlow v2.5.0。在 RTX 3090 上的 64 個data/config_oneblob.json實線）和 128 個（ benchmarks/bench_tensorflow.py ）神經benchmarks/bench_ours.cu寬多層感知器上測量。

用法

微型 CUDA 神經網路有一個簡單的 C++/CUDA API： 

# include < tiny-cuda-nn/common.h >

// Configure the model
nlohmann::json config = {
	{ " loss " , {
		{ " otype " , " L2 " }
	}},
	{ " optimizer " , {
		{ " otype " , " Adam " },
		{ " learning_rate " , 1e-3 },
	}},
	{ " encoding " , {
		{ " otype " , " HashGrid " },
		{ " n_levels " , 16 },
		{ " n_features_per_level " , 2 },
		{ " log2_hashmap_size " , 19 },
		{ " base_resolution " , 16 },
		{ " per_level_scale " , 2.0 },
	}},
	{ " network " , {
		{ " otype " , " FullyFusedMLP " },
		{ " activation " , " ReLU " },
		{ " output_activation " , " None " },
		{ " n_neurons " , 64 },
		{ " n_hidden_layers " , 2 },
	}},
};

using namespace tcnn ;

auto model = create_from_config(n_input_dims, n_output_dims, config);

// Train the model (batch_size must be a multiple of tcnn::BATCH_SIZE_GRANULARITY)
GPUMatrix< float > training_batch_inputs (n_input_dims, batch_size);
GPUMatrix< float > training_batch_targets (n_output_dims, batch_size);

for ( int i = 0 ; i < n_training_steps; ++i) {
	generate_training_batch (&training_batch_inputs, &training_batch_targets); // <-- your code

	float loss;
	model. trainer -> training_step (training_batch_inputs, training_batch_targets, &loss);
	std::cout << " iteration= " << i << " loss= " << loss << std::endl;
}

// Use the model
GPUMatrix< float > inference_inputs (n_input_dims, batch_size);
generate_inputs (&inference_inputs); // <-- your code

GPUMatrix< float > inference_outputs (n_output_dims, batch_size);
model.network-> inference (inference_inputs, inference_outputs);

範例：學習 2D 影像

我們提供了一個範例應用程序，其中學習了圖像函數(x,y) -> (R,G,B) 。它可以透過運行

tiny-cuda-nn$ ./build/mlp_learning_an_image data/images/albert.jpg data/config_hash.json

每幾個訓練步驟產生一個影像。在 RTX 4090 的預設配置下，每 1000 步驟應該花費 1 秒多一點的時間。

10步100步1000 步參考圖片
10步100步1000步參考

要求

  • NVIDIA GPU ；張量核心可以提高效能（如果可用）。顯示的所有結果均來自 RTX 3090。
  • 支援C++14 的編譯器。建議使用以下選項並經過測試：
    • Windows： Visual Studio 2019 或 2022
    • Linux： GCC/G++ 8 或更高版本
  • CUDA的最新版本。建議使用以下選項並經過測試：
    • Windows： CUDA 11.5 或更高版本
    • Linux： CUDA 10.2 或更高版本
  • CMake v3.21 或更高版本。
  • 此框架的完全融合的 MLP 元件在其預設配置中需要大量共享記憶體。它可能僅適用於 RTX 3090、RTX 2080 Ti 或更高階的 GPU。低端卡必須減少n_neurons參數或使用CutlassMLP （相容性更好但速度較慢）。

如果您使用的是 Linux，請安裝以下軟體包

sudo apt-get install build-essential git

我們也建議將 CUDA 安裝在/usr/local/中，並將 CUDA 安裝新增至您的 PATH 中。例如，如果您有 CUDA 11.4，請將以下內容新增至~/.bashrc 

 export PATH= " /usr/local/cuda-11.4/bin: $PATH "
export LD_LIBRARY_PATH= " /usr/local/cuda-11.4/lib64: $LD_LIBRARY_PATH "

編譯（Windows 和 Linux）

首先使用以下命令克隆此存儲庫及其所有子模組： 

$ git clone --recursive https://github.com/nvlabs/tiny-cuda-nn
$ cd tiny-cuda-nn

然後，使用 CMake 建置專案：（在 Windows 上，這必須在開發人員命令提示字元中） 

tiny-cuda-nn$ cmake . -B build -DCMAKE_BUILD_TYPE=RelWithDebInfo
tiny-cuda-nn$ cmake --build build --config RelWithDebInfo -j

如果編譯莫名其妙地失敗或花費超過一個小時，則可能會耗盡記憶體。在這種情況下，請嘗試在不帶-j的情況下執行上述命令。

PyTorch 擴展

tiny-cuda-nn附帶 PyTorch 擴展，允許在 Python 上下文中使用快速 MLP 和輸入編碼。這些綁定比完整的 Python 實作要快得多；特別是對於多重解析度雜湊編碼。

如果批次大小很小，Python/PyTorch 的開銷仍然會很大。例如，當批次大小為 64k 時，透過 PyTorch 捆綁的mlp_learning_an_image範例比原生 CUDA慢約 2 倍。當批量大小為 256k 或更高（預設）時，效能更加接近。

首先使用最新的、支援 CUDA 的 PyTorch 版本設定 Python 3.X 環境。然後，調用

pip install git+https://github.com/NVlabs/tiny-cuda-nn/ # subdirectory=bindings/torch

或者，如果您想從tiny-cuda-nn的本機複製進行安裝，請呼叫

tiny-cuda-nn$ cd bindings/torch
tiny-cuda-nn/bindings/torch$ python setup.py install

成功後，您可以使用tiny-cuda-nn模型，如下例所示： 

 import commentjson as json
import tinycudann as tcnn
import torch

with open ( "data/config_hash.json" ) as f :
	config = json . load ( f )

# Option 1: efficient Encoding+Network combo.
model = tcnn . NetworkWithInputEncoding (
	n_input_dims , n_output_dims ,
	config [ "encoding" ], config [ "network" ]
)

# Option 2: separate modules. Slower but more flexible.
encoding = tcnn . Encoding ( n_input_dims , config [ "encoding" ])
network = tcnn . Network ( encoding . n_output_dims , n_output_dims , config [ "network" ])
model = torch . nn . Sequential ( encoding , network )

有關範例，請參閱samples/mlp_learning_an_image_pytorch.py​​ 。

成分

以下是該框架的組成部分的摘要。 JSON 文件列出了設定選項。

網路
完全融合MLP src/fully_fused_mlp.cu快速實現小型多層感知器 (MLP)。
彎刀MLP src/cutlass_mlp.cu基於 CUTLASS 的 GEMM 例程的 MLP。比完全融合慢，但處理更大的網路並且仍然相當快。
輸入編碼
合成的include/tiny-cuda-nn/encodings/composite.h允許組合多種編碼。例如，可以用於組裝神經輻射快取編碼 [Müller 等人。 2021]。
頻率include/tiny-cuda-nn/encodings/frequency.h NeRF 的 [Mildenhall 等人。 2020]位置編碼同樣適用於所有維度。
網格include/tiny-cuda-nn/encodings/grid.h基於可訓練的多重解析度網格的編碼。用於即時神經圖形基元 [Müller 等人。 2022]。網格可以由哈希表、密集儲存或平鋪儲存支援。
身分include/tiny-cuda-nn/encodings/identity.h保持價值觀不變。
一個斑點include/tiny-cuda-nn/encodings/oneblob.h來自神經重要性採樣 [Müller 等人。 2019] 和神經控制變數 [Müller 等人。 2020]。
球諧函數include/tiny-cuda-nn/encodings/spherical_harmonics.h一種頻率空間編碼，比分量方式更適合方向向量。
三角波include/tiny-cuda-nn/encodings/triangle_wave.h NeRF 編碼的低成本替代方案。用於神經輻射緩存 [Müller 等人。 2021]。
損失
L1 include/tiny-cuda-nn/losses/l1.h標準 L1 損耗。
相對L1 include/tiny-cuda-nn/losses/l1.h透過網路預測標準化的相對 L1 損失。
MAPE include/tiny-cuda-nn/losses/mape.h平均絕對百分比誤差 (MAPE)。與相對 L1 相同，但由目標標準化。
SMAPE include/tiny-cuda-nn/losses/smape.h對稱平均絕對百分比誤差 (SMAPE)。與相對 L1 相同，但透過預測和目標的平均值進行歸一化。
L2 include/tiny-cuda-nn/losses/l2.h標準 L2 損耗。
相對L2 include/tiny-cuda-nn/losses/relative_l2.h透過網路預測標準化相對 L2 損失 [Lehtinen 等人。 2018]。
相對 L2 亮度include/tiny-cuda-nn/losses/relative_l2_luminance.h與上面相同，但透過網路預測的亮度進行歸一化。僅當網路預測為 RGB 時適用。用於神經輻射緩存 [Müller 等人。 2021]。
交叉熵include/tiny-cuda-nn/losses/cross_entropy.h標準交叉熵損失。僅當網路預測為 PDF 時適用。
變異數include/tiny-cuda-nn/losses/variance_is.h標準方差損失。僅當網路預測為 PDF 時適用。
最佳化器
亞當include/tiny-cuda-nn/optimizers/adam.h Adam 的實作 [Kingma 和 Ba 2014]，推廣到 AdaBound [Luo 等人。 2019]。
諾沃格勒include/tiny-cuda-nn/optimizers/lookahead.h Novograd 的實施 [Ginsburg 等。 2019]。
新加坡元include/tiny-cuda-nn/optimizers/sgd.h標準隨機梯度下降 (SGD)。
洗髮精include/tiny-cuda-nn/optimizers/shampoo.h二階洗髮精優化器的實現 [Gupta 等人。 2018] 進行了本土最佳化以及 Anil 等人的最佳化。 [2020]。
平均的include/tiny-cuda-nn/optimizers/average.h包裝另一個最佳化器並計算最後 N 次迭代中權重的線性平均值。平均值僅用於推理（不會反饋到訓練中）。
大量include/tiny-cuda-nn/optimizers/batched.h包裝另一個優化器，在平均梯度上每 N 步調用一次嵌套優化器。與增加批次大小具有相同的效果，但僅需要恆定的記憶體量。
合成的include/tiny-cuda-nn/optimizers/composite.h允許對不同參數使用多個優化器。
EMA include/tiny-cuda-nn/optimizers/average.h包裝另一個最佳化器並計算權重的指數移動平均值。平均值僅用於推理（不會反饋到訓練中）。
指數衰減include/tiny-cuda-nn/optimizers/exponential_decay.h包裝另一個最佳化器並執行分段恆定指數學習率衰減。
前瞻include/tiny-cuda-nn/optimizers/lookahead.h包裝另一個優化器，實現前瞻性演算法 [Zhang 等人。 2019]。

授權和引用

該框架根據 BSD 3-clause 許可證獲得許可。詳細資訊請參閱LICENSE.txt 。

如果您在研究中使用它，我們將不勝感激

 @software { tiny-cuda-nn ,
	author = { M"uller, Thomas } ,
	license = { BSD-3-Clause } ,
	month = { 4 } ,
	title = { {tiny-cuda-nn} } ,
	url = { https://github.com/NVlabs/tiny-cuda-nn } ,
	version = { 1.7 } ,
	year = { 2021 }
}

如需業務諮詢，請造訪我們的網站並提交表格：NVIDIA 研究許可

出版品和軟體

除此之外，該框架為以下出版物提供支援：

具有多重解析度雜湊編碼的即時神經圖形基元
湯瑪斯·穆勒、亞歷克斯·埃文斯、克里斯托夫·席德、亞歷山大·凱勒
ACM 圖形交易 ( SIGGRAPH )，2022 年 7 月
網站/論文/程式碼/影片/BibTeX

從影像中擷取三角形 3D 模型、材質和光照
雅各·芒克伯格、喬恩·哈塞爾格倫、沈天長、高軍、陳文正、亞歷克斯·埃文斯、托馬斯·穆勒、Sanja Fidler
CVPR（口頭） ，2022 年 6 月
網址/論文/影片/BibTeX

用於路徑追蹤的即時神經輻射緩存
湯瑪斯·穆勒、法布里斯·魯塞爾、揚·諾瓦克、亞歷山大·凱勒
ACM 圖形交易 ( SIGGRAPH )，2021 年 8 月
論文 / GTC 演講 / 影片 / 互動式結果檢視器 / BibTeX

以及以下軟體：

NerfAcc：通用 NeRF 加速工具箱
李瑞龍、Matthew Tancik、Angjoo Kanazawa
https://github.com/KAIR-BAIR/nerfacc

Nerfstudio：神經輻射場開發框架
Matthew Tancik*、Ethan Weber*、Evonne Ng*、李瑞龍、Brent Yi、Terrance Wang、Alexander Kristoffersen、Jake Austin、Kamyar Salahi、Abhik Ahuja、David McAllister、Angjoo Kanazawa
https://github.com/nerfstudio-project/nerfstudio

如果您的出版品或軟體未列出，請隨時提出拉取請求。

致謝

特別感謝 NRC 作者的有益討論，感謝 Nikolaus Binder 提供該框架的部分基礎設施，以及在利用 CUDA 內的 TensorCore 方面提供的幫助。

展開
附加信息
相關應用
爲您推薦
相關資訊 全部