tiny-cuda-nn
A neural network framework trains and runs small models locally on NVIDIA GPUs, with C++/CUDA and PyTorch interfaces for Windows and Linux.
tiny-cuda-nn is a neural network framework for developers and researchers building small models, NeRF renderers and neural graphics applications. It handles training and inference on local NVIDIA GPUs.
- Small neural networks: A fully fused multilayer perceptron (MLP) implementation targets small networks, while the slower CUTLASS implementation accommodates larger networks and a broader range of GPUs.
- Input encodings: Multiresolution hash encoding and NeRF positional encoding help developers represent inputs for neural graphics models, with support for combining encodings.
- PyTorch integration: Python projects can use the framework's networks and encodings through a PyTorch extension, which automatically uses half precision on GPUs with efficient support.
- Optional JIT fusion: Training and inference can use fusion for additional speed, and developers can integrate models into larger CUDA kernels, as Instant NGP does for NeRF rendering.
The framework runs on Windows and Linux and requires an NVIDIA GPU and CUDA. Tensor cores can improve performance. The default fully fused MLP needs substantial GPU shared memory; the project cites RTX 3090 and RTX 2080 Ti class hardware as likely suitable. Lower-end cards need smaller networks or the CUTLASS implementation.
Performance depends on the workload. Python and PyTorch overhead matters more at small batch sizes, and JIT fusion can slow training with large hash grids, wider networks or older GPUs.