tiny-cuda-nn - 专注于快速训练和查询神经网络的开源框架

Tiny CUDA Neural Networks 项目介绍

项目概述

Tiny CUDA Neural Networks 是一个小型、独立的框架，用于训练和查询神经网络。此框架的亮点在于其超快的“完全融合”多层感知机（MLP）实现，以及丰富的多分辨率哈希编码支持。除了这些，框架还支持多种输入编码、损失函数和优化器。

性能表现

该框架的多层感知机（MLP）在性能上表现卓越，尤其是在 NVIDIA RTX 3090 显卡上，比起 TensorFlow v2.5.0（使用 XLA 加速），其速度快得多。

使用方法

用户可以通过简单的 C++/CUDA API 来使用 Tiny CUDA Neural Networks。配置模型后，使用框架提供的函数训练模型，并在训练完毕后进行模型推理。示例如下：

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

// 配置模型
nlohmann::json config = {
    ...
};

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

// 训练模型
for (int i = 0; i < n_training_steps; ++i) {
    ...
    model.trainer->training_step(training_batch_inputs, training_batch_targets, &loss);
}

// 使用模型
GPUMatrix<float> inference_outputs(n_output_dims, batch_size);
model.network->inference(inference_inputs, inference_outputs);

示例: 学习一个二维图像

提供了一个示例应用程序，用于学习图像函数 (x,y) -> (R,G,B)。对于默认配置，每 1000 步大约只需 1 秒即可在 RTX 4090 上完成训练。

环境需求

需要 NVIDIA GPU，推荐使用支持 tensor cores 的高性能显卡，如 RTX 3090。
需要支持 C++14 的编译器，如 Windows 下的 Visual Studio 2019 或 Linux 下的 GCC/G++ 8 及更高版本。
需要安装 CUDA 工具包和 CMake。
示例可用于 PyTorch 环境，通过扩展提供快速的 MLP 和输入编码功能。

编译指南

克隆代码仓库：

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

使用 CMake 构建项目：

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

如果需要更长时间编译，可能由于内存不足，请去掉 -j 选项重试。

PyTorch 扩展

Tiny CUDA Neural Networks 提供了 PyTorch 扩展，允许在 Python 环境中使用其快速 MLP 和输入编码功能。其性能较于完整的 Python 实现要快，尤其是在多分辨率哈希编码时。

安装步骤：

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

或在本地克隆环境安装：

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

组件介绍

以下是框架中的主要组件：

网络

完全融合的 MLP：快速实现小型多层感知机。
CUTLASS MLP：基于 CUTLASS 的 GEMM 例程实现，适用于更大规模网络。

输入编码

复合编码、频率编码、网格编码等，支持多种数据编码需求。

损失函数

支持 L1、L2 等多种损失，以及各种相对损失和交叉熵损失。

优化器

提供 Adam、Novograd、SGD 等多种优化选项。

软件与出版物

Tiny CUDA Neural Networks 已在多篇学术论文中使用，相关研究覆盖了渲染、图像提取等多个领域。

许可证与引用

框架使用 BSD 3-clause 许可证。科研工作中如需引用，请参考以下引用格式：

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

以上是 Tiny CUDA Neural Networks 项目的简要介绍，涵盖了其功能、使用方法、组件及其在科研领域的应用。