tiny-cuda-nn

The tiny-cuda-nn framework consists of several specialized components for neural network training and inference. These include highly optimized networks, various input encodings, loss functions, and optimizers.

Networks

• Fully fused MLP: A lightning-fast implementation for small multi-layer perceptrons (MLPs).
• CUTLASS MLP: An MLP based on CUTLASS GEMM routines. It is slower than the fully-fused version but handles larger networks.

Input Encodings

• Composite: Composes multiple encodings (e.g., for Neural Radiance Caching).
• Frequency: NeRF-style positional encoding applied equally to all dimensions.
• Grid: Trainable multiresolution grids (used in Instant NGP). Supports hashtables, dense storage, or tiled storage.
• Identity: Leaves input values untouched.
• Oneblob: Based on Neural Importance Sampling and Neural Control Variates.
• SphericalHarmonics: Frequency-space encoding suitable for direction vectors.
• TriangleWave: A low-cost alternative to NeRF's encoding.

Losses

• L1 / Relative L1: Standard and prediction-normalized L1 loss.
• MAPE / SMAPE: Mean absolute percentage error and symmetric MAPE.
• L2 / Relative L2 / Relative L2 Luminance: Standard L2, prediction-normalized L2, and luminance-normalized L2 (for RGB predictions).
• Cross Entropy: Standard cross entropy (for PDF predictions).
• Variance: Standard variance loss (for PDF predictions).

Optimizers

• Adam: Implementation of Adam/AdaBound.
• Novograd: Implementation of Novograd.
• SGD: Standard stochastic gradient descent.
• Shampoo: 2nd order Shampoo optimizer.
• Average / EMA: Wrappers that compute a linear or exponential moving average of weights for inference.
• Batched: Wraps an optimizer to invoke it every N steps on averaged gradients (simulates larger batch size with constant memory).
• Composite: Allows using different optimizers on different parameters.
• Exponential Decay: Performs piecewise-constant exponential learning-rate decay.
• Lookahead: Implements the lookahead algorithm.

JIT fusion (available in v2.0+) provides a 1.5x to 2.5x performance boost by compiling models into CUDA device functions via Runtime Compilation (RTC). It is highly recommended to enable it.

C++ Usage: Set the jit_fusion property using set_jit_fusion(tcnn::supports_jit_fusion()).

Python Usage: Set the jit_fusion property on the model instance.

Note: If JIT compilation fails, a warning is emitted and the system automatically falls back to the standard 1.X code path. For very large models (20M+ parameters) or older GPUs (RTX 3000 series or earlier), JIT fusion might slow down training; test separately for training and inference in those cases.

// C++
auto model = tcnn::create_from_config(...);
model->set_jit_fusion(tcnn::supports_jit_fusion());

# Python
import tinycudann as tcnn
model = tcnn.NetworkWithInputEncoding(...) 
model.jit_fusion = tcnn.supports_jit_fusion()

For maximum performance (e.g., fusing a model into a ray marcher), you can manually integrate a tiny-cuda-nn model into a custom CUDA kernel using the CudaRtcKernel API.

Steps:

1. Convert your kernel into a string.
2. Prepend the model's device function using model->generate_device_function("function_name") via the {MODEL_DEVICE_FUNCTION} placeholder.
3. Use tcnn::CudaRtcKernel to compile and launch the fused kernel.

Requirement: All 32 threads of the warp must be active when calling the model function.

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

auto model = tcnn::create_from_config(32 /* input dims */, 16 /* output dims */, ...);
auto fused_kernel = tcnn::CudaRtcKernel(
    "your_kernel",
    fmt::format(R"(
        {MODEL_DEVICE_FUNCTION}
        __global__ void your_kernel(...) {
            // Get input to model from either registers or memory.
            tcnn::hvec<32> input = ...;
            // Call tiny-cuda-nn model. All 32 threads of the warp must be active here.
            tcnn::hvec<16> output = model_fun(nerf_in, params); 
            // Do something with the model output.
        }",
        fmt::arg("MODEL_DEVICE_FUNCTION", model->generate_device_function("model_fun")),
    )
);

uint32_t blocks = 1;
uint32_t threads = 128; // Must be multiple of 32 for neural networks to work.
uint32_t shmem_size = 0;
cudaStream_t stream = nullptr;
fused_kernel.launch(blocks, threads, shmem_size, stream, ... /* params of your_kernel */);

Tiny CUDA neural networks provide a C++/CUDA API for configuring, training, and performing inference with models like Fully Fused MLPs and Multiresolution Hash Encodings. Models are configured using nlohmann::json objects.

Key steps:

1. Configure: Define loss, optimizer, encoding, and network parameters in a JSON object.
2. Create: Use tcnn::create_from_config(n_input_dims, n_output_dims, config) to instantiate the model.
3. Train: Use model.trainer->training_step(inputs, targets, &loss) within a training loop. Note that batch_size must be a multiple of tcnn::BATCH_SIZE_GRANULARITY.
4. Inference: Use model.network->inference(inputs, outputs) to query the model.

#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);
model->set_jit_fusion(supports_jit_fusion()); // Optional: accelerate with JIT fusion

// 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);

Tiny CUDA NN provides several wrapper optimizers that can be nested within a nested key to modify the behavior of a base optimizer:

• Average: Computes a linear average of parameters over the last n_samples steps. Used for inference only.
• Batched: Invokes the nested optimizer once every batch_size_multiplier steps on the averaged gradient. Mimics larger batch sizes with constant memory.
• EMA: Computes an exponential moving average of parameters using a decay factor. Used for inference only.
• ExponentialDecay: Performs piecewise-constant exponential learning-rate decay between decay_start and decay_end with a specific decay_interval and decay_base.
• Lookahead: Implements the lookahead algorithm using alpha (distance fraction) and n_steps (nested steps per lookahead step).

// Example: Lookahead wrapping Adam
{
	"otype": "Lookahead",
	"alpha": 0.5,
	"n_steps": 16,
	"nested": {
		"otype": "Adam"
	}
}