微型 CUDA 神经网络

这是一个小型、自包含的框架，用于训练和查询神经网络。最引人注目的是，它包含一个极快的“完全融合”多层感知机（技术论文），一种多功能的多分辨率哈希编码（技术论文），以及对各种其他输入编码、损失函数和优化器的支持。

性能

完全融合网络与 TensorFlow v2.5.0 w/ XLA 的对比。在 RTX 3090 上，分别测量了宽度为 64（实线）和 128（虚线）个神经元的多层感知机。由 benchmarks/bench_ours.cu 和 benchmarks/bench_tensorflow.py 使用 data/config_oneblob.json 生成。

使用方法

微型 CUDA 神经网络提供了一个简单的 C++/CUDA API：

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

// 配置模型
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()); // 可选：通过 JIT 融合加速

// 训练模型（batch_size 必须是 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); // <-- 你的代码

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

// 使用模型
GPUMatrix<float> inference_inputs(n_input_dims, batch_size);
generate_inputs(&inference_inputs); // <-- 你的代码

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

JIT 融合

JIT 融合是 tiny-cuda-nn v2.0 及更高版本中的一项新功能，属于可选特性。 根据模型和 GPU 的不同，几乎总是建议启用自动 JIT 融合，以获得 1.5 到 2.5 倍的性能提升。 较新的 GPU 通常会带来更大的加速效果。

如果您的模型包含非常大的哈希网格（约 2000 万+ 参数）或 MLP（每层神经元数量超过 128 个），或者您的 GPU 是 RTX 3000 系列及更早型号，则 JIT 融合可能会减慢训练速度。 在极少数情况下，推理也会变慢。 在这种情况下，建议分别尝试为训练和推理启用 JIT 融合，以衡量是否确实更快。

如果您在其他情况下遇到性能下降，或在启用 JIT 融合时遇到其他问题，请提交一个问题。

自动 JIT 融合

要启用 JIT 融合，只需将模型的 jit_fusion 属性设置为 true。 此后，无论进行推理还是训练，模型都将使用 JIT 模式。 请注意，如果 JIT 编译过程中出现错误，系统将发出警告，并自动关闭 JIT 编译模式。 此时，您的代码仍将使用 tiny-cuda-nn 1.X 的代码路径运行。

auto model = tcnn::create_from_config(...);
model->set_jit_fusion(tcnn::supports_jit_fusion()); // 如果系统支持，则启用 JIT

JIT 融合也可以通过 PyTorch 绑定来启用，但加速效果会较低，尤其是在训练阶段。 这是因为 JIT 编译器无法访问完整的计算图，因此能够融合和优化的内容较少。

import tinycudann as tcnn

model = tcnn.NetworkWithInputEncoding(...) # 或任何其他 tcnn 模型
model.jit_fusion = tcnn.supports_jit_fusion() # 如果系统支持，则启用 JIT

手动 JIT 融合

当应用程序与 JIT 融合更紧密地集成时，可以获得更大的加速效果。 例如，Instant NGP 通过将整个 NeRF 射线追踪器融合到一个内核中，实现了 5 倍的加速。

JIT 融合的工作原理是将给定的 tiny-cuda-nn 模型转换为 CUDA 设备函数，然后利用 CUDA 的运行时编译 (RTC) 功能将其编译成一个内核。

要将 tiny-cuda-nn 模型与您应用中的更大内核集成，您需要：

1. 将您的内核转换为字符串，
2. 在其前缀添加 tiny-cuda-nn 模型的设备函数，
3. 将结果传递给 tiny-cuda-nn 的运行时编译 API。

以下是一个示例，展示了如何使用具有 32 个输入维度和 16 个输出维度的 tiny-cuda-nn 模型实现一个最小内核：

#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(...) {
            // 从寄存器或内存中获取模型输入。
            tcnn::hvec<32> input = ...;
            // 调用 tiny-cuda-nn 模型。在此处，warp 中的所有 32 个线程都必须处于活动状态。
            tcnn::hvec<16> output = model_fun(nerf_in, params); 
            // 对模型输出做些处理。
        }",
        fmt::arg("MODEL_DEVICE_FUNCTION", model->generate_device_function("model_fun")),
    )
);

uint32_t blocks = 1;
uint32_t threads = 128; // 必须是 32 的倍数，以便神经网络正常工作。
uint32_t shmem_size = 0; // 可以根据 your_kernel 的需求设置任意大小。
cudaStream_t stream = nullptr; // 可以使用任意流。
fused_kernel.launch(blocks, threads, shmem_size, stream, ... /* your_kernel 的参数 */);

以下是 Instant NGP 的 NeRF 集成与 JIT 编译器的参考：

• src/testbed_nerf.cu
• include/neural-graphics-primitives/fused_kernels/render_nerf.cuh

示例：学习一张2D图像

我们提供了一个示例应用，用于学习一个图像函数 (x,y) -> (R,G,B)。可以通过以下命令运行：

tiny-cuda-nn$ ./build/mlp_learning_an_image https://oss.gittoolsai.com/images/NVlabs_tiny-cuda-nn_readme_974957d202f6.jpg data/config_hash.json

该程序会在每隔几个训练步骤生成一张图像。在 RTX 4090 上使用默认配置时，每 1000 步大约需要 1 秒多一点。

10 步	100 步	1000 步	参考图像

要求

• 一块 NVIDIA GPU；如果支持张量核心，则可以进一步提升性能。所有展示的结果均来自 RTX 3090。
• 一个支持 C++17 的编译器。推荐并经过测试的选项如下：
  • 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 12.6.3，可以在 ~/.bashrc 文件中添加以下内容：

export PATH="/usr/local/cuda-12.6.3/bin:$PATH"
export LD_LIBRARY_PATH="/usr/local/cuda-12.6.3/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 时，捆绑的 mlp_learning_an_image 示例通过 PyTorch 运行的速度比原生 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

默认情况下，该扩展会在支持半精度计算的 GPU（Volta、Turing、Ampere 等）上自动启用 FP16，在旧架构或 FP16 性能较差的硬件（如 Pascal/GTX 10 系列）上禁用 FP16。

如果你想覆盖此行为（例如强制在不支持的硬件上启用 FP16，或为了调试而禁用 FP16），可以在安装前设置 TCNN_HALF_PRECISION 环境变量：

禁用 FP16：0
启用 FP16：1

示例：

# Linux / macOS（禁用 FP16）
export TCNN_HALF_PRECISION=0
pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch

安装成功后，你可以按照以下示例使用 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)

# 选项 1：高效的编码+网络组合。
model = tcnn.NetworkWithInputEncoding(
	n_input_dims, n_output_dims,
	config["encoding"], config["network"]
)

# 选项 2：分离模块。速度稍慢但更灵活。
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)

model.jit_fusion = tcnn.supports_jit_fusion() # 可选：通过 JIT 融合加速

更多示例请参阅 samples/mlp_learning_an_image_pytorch.py。

组件

以下是该框架的组件概览。JSON 文档列出了配置选项。

网络
全融合 MLP	src/fully_fused_mlp.cu	超快速的小型多层感知机（MLP）实现。
CUTLASS 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	不对输入值进行任何变换。
Oneblob 编码	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 的情况。

优化器
Adam	include/tiny-cuda-nn/optimizers/adam.h	实现了 Adam [Kingma 和 Ba, 2014]，并扩展为 AdaBound [Luo 等人, 2019]。
Novograd	include/tiny-cuda-nn/optimizers/lookahead.h	实现了 Novograd [Ginsburg 等人, 2019]。
SGD	include/tiny-cuda-nn/optimizers/sgd.h	标准随机梯度下降（SGD）。
Shampoo	include/tiny-cuda-nn/optimizers/shampoo.h	实现了二阶 Shampoo 优化器 [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	包装另一个优化器，执行分段常数的指数学习率衰减。
Lookahead 优化器	include/tiny-cuda-nn/optimizers/lookahead.h	包装另一个优化器，实现了 lookahead 算法 [Zhang 等人, 2019]。

许可与引用

本框架采用 BSD 3 条款许可证授权。详情请参阅 LICENSE.txt 文件。

如果您在研究中使用本框架，请通过以下 BibTeX 格式引用：

@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 = {2.0},
	year = {2021}
}

如需商业合作，请访问我们的官网并提交表格：NVIDIA Research Licensing

出版物与软件

该框架支持以下出版物：

具有多分辨率哈希编码的即时神经图形基元
托马斯·穆勒、亚历克斯·埃文斯、克里斯托夫·希德、亚历山大·凯勒
ACM 图形学汇刊 (SIGGRAPH)，2022年7月
网站 / 论文 / 代码 / 视频 / BibTeX

从图像中提取三角形3D模型、材质和光照
雅各布·蒙克贝格、乔恩·哈塞尔格伦、沈天畅、高俊、陈文政、亚历克斯·埃文斯、托马斯·穆勒、桑雅·菲德勒
CVPR（口头报告），2022年6月
网站 / 论文 / 视频 / BibTeX

用于路径追踪的实时神经辐射缓存
托马斯·穆勒、法布里斯·鲁塞尔、扬·诺瓦克、亚历山大·凯勒
ACM 图形学汇刊 (SIGGRAPH)，2021年8月
论文 / GTC演讲 / 视频 / 交互式结果查看器 / BibTeX

此外，该框架还支持以下软件：

NerfAcc：通用NeRF加速工具箱
李瑞龙、马修·坦西克、安久·卡纳扎瓦
https://github.com/KAIR-BAIR/nerfacc

Nerfstudio：神经辐射场开发框架
马修·坦西克*、伊森·韦伯*、艾沃妮·吴*、李瑞龙、布伦特·易、泰伦斯·王、亚历山大·克里斯托弗森、杰克·奥斯汀、卡米亚尔·萨拉希、阿比克·阿胡贾、大卫·麦卡利斯特、安久·卡纳扎瓦
https://github.com/nerfstudio-project/nerfstudio

如果您发现自己的出版物或软件未在此列出，欢迎随时提交Pull Request。

致谢

特别感谢NRC的作者们提供的有益讨论，以及尼古劳斯·宾德为本框架提供部分基础设施，并协助在CUDA中利用TensorCore。

Tiny CUDA Neural Networks 快速上手指南

Tiny CUDA Neural Networks (tcnn) 是一个轻量级、自包含的框架，专为在 NVIDIA GPU 上高效训练和推理神经网络而设计。其核心优势在于提供了极速的“全融合”多层感知机（Fully Fused MLP）和多分辨率哈希编码（Multiresolution Hash Encoding），广泛应用于即时神经图形基元（如 Instant NGP）等高性能场景。

环境准备

在开始之前，请确保您的系统满足以下硬件和软件要求：

硬件要求

• GPU: 必须配备 NVIDIA GPU。若具备 Tensor Cores（如 RTX 20/30/40 系列或 A100/H100），性能将显著提升。
  • 注意：默认的“全融合”MLP 需要大量共享内存，建议在 RTX 3090、RTX 2080 Ti 或更高阶显卡上运行。低端显卡需减少 n_neurons 参数或改用 CutlassMLP。

软件依赖

• 操作系统: Windows 或 Linux
• 编译器: 支持 C++17
  • Windows: Visual Studio 2019 或 2022
  • Linux: GCC/G++ 8 或更高版本
• CUDA Toolkit:
  • Windows: CUDA 11.5 或更高
  • Linux: CUDA 10.2 或更高
• 构建工具: CMake v3.21 或更高
• 其他 (Linux): build-essential, git

Linux 用户前置安装命令：

sudo apt-get install build-essential git

配置 CUDA 环境变量 (Linux 示例): 建议将 CUDA 安装在 /usr/local/ 并添加到路径。以 CUDA 12.6.3 为例，请在 ~/.bashrc 中添加：

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

安装步骤

方式一：编译 C++ 原生库

1. 克隆仓库 使用 --recursive 参数拉取代码及子模块：

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

2. 构建项目 使用 CMake 进行配置和编译（Windows 用户请在 Developer Command Prompt 中执行）：

   cmake . -B build -DCMAKE_BUILD_TYPE=RelWithDebInfo
   cmake --build build --config RelWithDebInfo -j
   

   提示：如果编译过程中内存不足导致失败或耗时过长，请移除 -j 参数尝试单线程编译。

方式二：安装 PyTorch 扩展 (Python)

如果您希望在 Python 环境中使用，请先确保已安装支持 CUDA 的 PyTorch。

在线安装：

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

本地源码安装：

cd bindings/torch
python setup.py install

关于半精度 (FP16) 的控制： 默认情况下，库会自动检测显卡架构启用或禁用 FP16。如需强制控制，可在安装前设置环境变量：

• 禁用 FP16: export TCNN_HALF_PRECISION=0
• 启用 FP16: export TCNN_HALF_PRECISION=1

基本使用

1. Python (PyTorch) 使用示例

这是最常用的方式，适合快速原型开发和集成到现有 PyTorch 项目中。

import commentjson as json
import tinycudann as tcnn
import torch

# 加载配置文件 (假设 config_hash.json 存在于当前目录)
with open("data/config_hash.json") as f:
    config = json.load(f)

# 定义输入输出维度
n_input_dims = 2
n_output_dims = 3

# 创建模型：组合编码器和网络
model = tcnn.NetworkWithInputEncoding(
    n_input_dims=n_input_dims,
    n_output_dims=n_output_dims,
    encoding_config=config["encoding"],
    network_config=config["network"]
)

# 可选：启用 JIT 融合以获得 1.5x - 2.5x 的性能提升
model.jit_fusion = tcnn.supports_jit_fusion()

# 准备数据 (需在 GPU 上)
inputs = torch.rand(1024, n_input_dims, device="cuda", dtype=torch.float32)
targets = torch.rand(1024, n_output_dims, device="cuda", dtype=torch.float32)

# 前向传播
outputs = model(inputs)

# 计算损失并反向传播
loss = ((outputs - targets) ** 2).mean()
loss.backward()

# 优化器步骤 (需配合 torch.optim 使用)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
optimizer.step()

2. C++ 原生 API 使用示例

适合对性能有极致追求或需要深度定制 CUDA Kernel 的场景。

#include <tiny-cuda-nn/common.h>
#include <tiny-cuda-nn/trainer.h>
#include <nlohmann/json.hpp>

using namespace tcnn;

int main() {
    // 1. 配置模型 (JSON 格式)
    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},
        }},
    };

    uint32_t n_input_dims = 2;
    uint32_t n_output_dims = 3;

    // 2. 创建模型
    auto model = create_from_config(n_input_dims, n_output_dims, config);
    
    // 可选：启用 JIT 融合加速
    model->set_jit_fusion(supports_jit_fusion());

    // 3. 准备训练数据 (GPUMatrix 需位于显存)
    uint32_t batch_size = 1024; 
    // 注意：batch_size 必须是 tcnn::BATCH_SIZE_GRANULARITY 的倍数
    GPUMatrix<float> training_batch_inputs(n_input_dims, batch_size);
    GPUMatrix<float> training_batch_targets(n_output_dims, batch_size);

    // 此处应填入生成/加载数据的逻辑
    // generate_training_batch(&training_batch_inputs, &training_batch_targets); 

    // 4. 训练循环
    int n_training_steps = 1000;
    for (int i = 0; i < n_training_steps; ++i) {
        float loss;
        // 执行一步训练
        model->trainer->training_step(training_batch_inputs, training_batch_targets, &loss);
        
        if (i % 100 == 0) {
            std::cout << "iteration=" << i << " loss=" << loss << std::endl;
        }
    }

    // 5. 推理
    GPUMatrix<float> inference_outputs(n_output_dims, batch_size);
    model->network->inference(training_batch_inputs, inference_outputs);

    return 0;
}

性能提示

• JIT Fusion: 在大多数现代 GPU (RTX 3000 系列及更新) 上，强烈建议开启 jit_fusion，可获得显著加速。若遇到显存不足或大模型训练变慢的情况，可尝试关闭它。
• Batch Size: PyTorch 绑定在小 Batch Size 下会有较大的 Python 开销，建议尽量使用较大的 Batch Size (如 64k 以上) 以发挥最大性能。