如果你经常接触深度学习,一定对 PyTorch 或 TensorFlow 的 API 很熟悉。tensor + tensor 就能做加法,.backward() 就能自动求梯度,nn.Module 就能搭网络------这些体验已经成了行业标配。但你是否想过,这些"魔法"背后到底发生了什么?
很多人选择调包是因为"业务需要",但如果你真的想理解深度学习框架的底层原理,最好的方式就是自己动手写一个。不是用 Python 包装一层 numpy,而是用 C++ 从零实现张量运算、自动微分、神经网络层、优化器,甚至分布式训练。
这个项目就是一个纯 C++17 实现的轻量级深度学习训练框架。它不依赖任何外部深度学习库(没有 CUDA 的 cuBLAS/cuDNN,没有 MKL),从零实现了 PyTorch 风格的 API:张量(Tensor)、自动微分(Autograd)、神经网络模块(nn.Module)、优化器(Optimizer)。同时支持 CPU(带 BLAS 加速)和 CUDA GPU,支持 FP16/FP32/BF16 混合精度,还支持多机多 GPU 的分布式训练(通过 NCCL)。
更重要的是,它的代码结构清晰、分层明确,是学习深度学习框架内部原理的绝佳材料。
自动微分技术科普:计算图、前向传播与反向传播
在深入代码之前,我们先搞清楚一个核心概念:自动微分(Automatic Differentiation,AD)。
很多人会把自动微分和符号微分、数值微分搞混。简单区分一下:
- 符号微分:对表达式做代数变换求出精确的导数公式。问题是表达式复杂时会出现"表达式膨胀"(expression swell),而且无法处理含控制流的代码。
- 数值微分:用差分近似 f'(x) ≈ (f(x+h) - f(x)) / h。简单但精度低(截断误差),且计算量大(每个参数都要单独算一次)。
- 自动微分:把程序拆解成基本运算的序列,利用链式法则逐节点计算梯度。既不产生表达式膨胀,也没有截断误差,是深度学习框架的标配方案。
自动微分的核心数据结构是计算图(Computation Graph)。每次对张量做运算时,框架会记录这个操作以及它的输入输出,形成一个有向无环图(DAG)。反向传播时,从损失节点开始,沿着图的边反向遍历,应用链式法则把梯度逐层传递回去。
画个简单的例子:
a=2.0 b=3.0
\ /
\ /
e = a * b (乘法)
|
v
d = e + c (加法, c=10.0)
|
v
L = d * f (乘法, f=-2.0) ← 损失节点
求 L 对 a 的梯度时,反向传播沿着 L <- d <- e <- a 这条路径走:
- dL/de = dL/dd * dd/de = f * 1 = -2.0
- dL/da = dL/de * de/da = f * b = -2.0 * 3.0 = -6.0
这就是链式法则的威力:不需要手动推导公式,框架自动帮你算。
项目架构总览:四层设计哲学
这个框架采用了清晰的分层架构,从底层到顶层分为四层:
+--------------------------------------------------+
| 应用层:训练/推理代码 |
| (examples/mnist, examples/ddp, 自定义模型) |
+--------------------------------------------------+
| Module 层 |
| (nn.Linear, nn.Conv2d, nn.BatchNorm, Loss) |
| 高层神经网络构建块,封装了参数管理和前向逻辑 |
+--------------------------------------------------+
| Function 层 |
| (FuncActivation, FuncElemWise, FuncLoss, ...) |
| 自动微分函数定义:forward 计算 + backward 梯度规则 |
+--------------------------------------------------+
| Operation 层 |
| (OpElemWise{Cpu,Cuda}, OpLinalg{Cpu,Cuda}, ...) |
| 实际计算内核:CPU 实现 + CUDA 实现 |
+--------------------------------------------------+
| Tensor 层 |
| (Tensor, Storage, Allocator, Device, DType) |
| 张量基础抽象:内存管理、设备抽象、类型系统 |
+--------------------------------------------------+
每一层都有明确的职责边界:
- Tensor 层:负责内存管理(自定义分配器 + 缓存策略)、设备抽象(CPU/GPU 统一接口)、数据类型(FP32/FP16/BF16)。这是整个框架的地基。
- Operation 层:实现具体的数学运算内核。每个运算都有 CPU 和 CUDA 两个版本,通过分发机制(Dispatch)自动选择。
- Function 层:定义自动微分的"函数节点"。每个函数节点包含 forward(前向计算)和 backward(反向梯度)两个方法。这是自动微分引擎的核心。
- Module 层:提供 nn.Module 风格的高层抽象,让开发者像搭积木一样构建神经网络。
这种分层设计的好处是:加一个新的运算只需要在 Operation 层加 CPU/CUDA 实现,在 Function 层定义反向规则,上层代码完全不用改。
Tensor 核心实现:内存管理与设备抽象
Tensor 是整个框架最核心的类。它的设计借鉴了 PyTorch 的 Tensor 抽象,但用 C++ 重新实现。
智能指针与 Pimpl 模式
cpp
Tensor 使用智能指针管理,避免内存泄漏。核心的 Tensor 类结构如下:
class Tensor {
public:
// 创建操作
static Tensor zeros(const std::vector<int64_t>& shape, const TensorOptions& options);
static Tensor ones(const std::vector<int64_t>& shape, const TensorOptions& options);
static Tensor randn(const std::vector<int64_t>& shape, const TensorOptions& options);
// 算术运算符重载(这也是构建计算图的关键)
Tensor operator+(const Tensor& other) const;
Tensor operator-(const Tensor& other) const;
Tensor operator*(const Tensor& other) const;
Tensor operator/(const Tensor& other) const;
Tensor operator^(const Tensor& exponent) const;
// 自动微分接口
void backward();
void zeroGrad();
Tensor grad() const;
// 内存与设备
const std::shared_ptr<TensorImpl>& impl() const;
Device device() const;
Tensor to(const Device& device);
private:
std::shared_ptr<TensorImpl> impl_; // Pimpl 指针
};
using TensorPtr = std::shared_ptr<Tensor>;
这里用了 Pimpl(Pointer to Implementation)模式:Tensor 类只持有一个 shared_ptr,具体的实现细节(数据指针、shape、dtype、device 等)都封装在 TensorImpl 里。这样做的好处是编译依赖少、二进制接口稳定。
内存分配器与缓存策略
框架实现了自定义的内存分配器(Allocator)和缓存分配器(CachedAllocator)。CUDA 版本会使用 CUDA 的内存池来减少设备内存分配的开销。
cpp
class Allocator {
public:
virtual ~Allocator() = default;
virtual void* allocate(size_t bytes) = 0;
virtual void deallocate(void* ptr, size_t bytes) = 0;
};
CPU 使用 std::malloc/free,CUDA 使用 cudaMalloc/cudaFree(带缓存策略)。CachedAllocator 会缓存已释放的内存块,避免频繁的 malloc/free 系统调用。对于 GPU 来说,内存分配是非常昂贵的操作(需要和驱动交互),所以缓存策略尤为重要。
设备抽象与类型系统
Device 枚举区分 CPU 和 CUDA:
cpp
enum class DeviceType { CPU, CUDA };
class Device {
public:
static Device cpu();
static Device cuda(int index = 0);
DeviceType type() const;
int index() const;
};
enum class DType { FLOAT32, FLOAT16, BFLOAT16, INT32, INT64 };
TensorOptions 把 device 和 dtype 打包在一起,创建张量时指定一次,后续所有操作都继承这个选项。
自动微分引擎:Function 节点与动态计算图
这是整个框架最精妙的部分。
Function 基类设计
每个运算都会创建一个 Function 节点,继承自抽象基类:
cpp
class Function {
public:
virtual ~Function() = default;
// 前向计算:接收输入张量,返回输出张量
virtual Tensor forward(const std::vector<Tensor>& inputs) = 0;
// 反向传播:接收输出梯度,返回各输入的梯度
// nextFunctions 记录了计算图的拓扑关系
std::vector<Function*> nextFunctions;
// 反向传播的钩子(用于调试或自定义梯度)
std::vector<Function::Hook> hooks;
};
每个具体的运算(如加法、乘法、激活函数)都实现自己的 Function 子类,重写 forward 和 backward 方法。backward 方法里实现的就是该运算的梯度公式。
运算符重载即建图
当你写 c = a * b 时,运算符重载不仅仅是做数学计算,还会:
- 创建一个新的 Function 节点(MulFunction)
- 把 a 和 b 的 nextFunctions 列表里添加这个新节点
- 执行实际的矩阵乘法运算
- 返回结果张量
这样就自动构建了一棵计算图。当你调用 c.backward() 时,框架从 c 节点开始,沿着 nextFunctions 反向遍历,依次调用每个 Function 的 backward 方法计算梯度。
NoGradGuard:RAII 风格的梯度控制
框架用了 C++ 的 RAII(Resource Acquisition Is Initialization)惯用法来实现梯度开关:
cpp
class NoGradGuard {
public:
NoGradGuard();
~NoGradGuard();
// 禁止拷贝
NoGradGuard(const NoGradGuard&) = delete;
NoGradGuard& operator=(const NoGradGuard&) = delete;
private:
bool prev_value_; // 保存之前的梯度开关状态
};
使用方式:
{
NoGradGuard guard; // 进入作用域,禁用梯度追踪
auto output = model.forward(input); // 这些操作不会记录到计算图中
// ...
} // 离开作用域,自动恢复之前的梯度状态
这比 Python 的 with torch.no_grad(): 更优雅,因为 C++ 的析构函数保证在任何路径(包括异常)下都能正确恢复状态。
运算层实现:CPU/GPU 双后端分发
Operation 层是实际干活的层。每个运算都有 CPU 和 CUDA 两个实现,通过分发机制自动选择。
分发机制(Dispatch)
Dispatch.h 根据 Tensor 的 device 类型,把运算分发到对应的 CPU 或 CUDA 实现:
cpp
template <typename Op>
Tensor dispatch(const Tensor& a, const Tensor& b) {
if (a.device().type() == DeviceType::CPU) {
return Op::cpu_impl(a, b);
} else if (a.device().type() == DeviceType::CUDA) {
return Op::cuda_impl(a, b);
}
}
CUDA 内核实现示例
以逐元素加法为例,CUDA 版本的 kernel 大致如下:
cpp
template <typename scalar_t>
__global__ void elemwise_add_kernel(
const scalar_t* __restrict__ a,
const scalar_t* __restrict__ b,
scalar_t* __restrict__ output,
int64_t total_elements) {
int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < total_elements) {
output[idx] = a[idx] + b[idx];
}
}
template <typename scalar_t>
void elemwise_add_cuda(const Tensor& a, const Tensor& b, Tensor& output) {
int64_t total = a.numel();
int block_size = 256;
int grid_size = (total + block_size - 1) / block_size;
elemwise_add_kernel<scalar_t><<<grid_size, block_size>>>(
a.data_ptr<scalar_t>(),
b.data_ptr<scalar_t>(),
output.data_ptr<scalar_t>(),
total
);
}
框架对每个基本运算都实现了 CPU(使用 BLAS 或手写循环)和 CUDA(使用 GPU kernel)两个版本,开发者不需要关心底层用的是什么设备。
融合操作(Fused Operations)
为了进一步性能优化,框架还支持操作融合(Operator Fusion)。比如 add + activation 可以融合成一个 kernel,避免中间结果的读写。这在 LLM 推理中特别重要,因为激活函数的数量远超矩阵乘法。
神经网络模块:nn.Module 设计模式
Module 层提供了 PyTorch 风格的 nn.Module 抽象,让构建网络变得简单直观。
Module 基类
cpp
class Module {
public:
virtual Tensor forward(Tensor input) = 0;
virtual Tensor forward(Tensor input, Tensor mask) {
return forward(input);
}
// 参数管理
std::vector<TensorPtr> parameters();
std::vector<TensorPtr> buffers();
// 设备迁移
void to(const Device& device);
void train(bool mode = true);
void eval();
virtual ~Module() = default;
};
常用层实现
框架实现了常用的神经网络层:
- Linear:全连接层 y = xW^T + b
- Conv2d:二维卷积(支持多种 padding/stride/dilation)
- BatchNorm2d:批归一化
- ReLU/Sigmoid/Tanh/GELU:各种激活函数
- CrossEntropyLoss/MSELoss:常用损失函数
- Dropout:随机失活
构建网络的示例
cpp
class SimpleNet : public Module {
public:
SimpleNet(int in_features, int hidden_features, int out_features) {
linear1 = std::make_shared<Linear>(in_features, hidden_features);
linear2 = std::make_shared<Linear>(hidden_features, out_features);
}
Tensor forward(Tensor input) override {
Tensor x = linear1->forward(input);
x = relu->forward(x);
x = linear2->forward(x);
return x;
}
private:
std::shared_ptr<Linear> linear1;
std::shared_ptr<Linear> linear2;
std::shared_ptr<ReLU> relu;
};
优化器实现:从 SGD 到 Adam 的统一接口
优化器是训练循环的核心。框架实现了一个统一的优化器基类,以及多种常用优化器的具体实现。
优化器基类设计
cpp
class Optimizer {
public:
Optimizer(std::vector<TensorPtr>& parameters, float lr, float weightDecay = 0.f);
virtual ~Optimizer() = default;
void step() {
NoGradGuard guard; // 参数更新不需要梯度追踪
step_++;
doStep();
}
virtual void doStep() = 0; // 各优化器实现自己的更新规则
float getLr() const { return lr_; }
void setLr(float lr) { this->lr_ = lr; }
void zeroGrad() const; // 清零梯度
protected:
virtual Tensor getDecayedGrad(TensorPtr param) const; // 权重衰减
void initCache(std::vector<Tensor>& cache, bool setZero = false);
std::vector<TensorPtr> parameters_;
int64_t step_;
float lr_;
float weightDecay_;
};
SGD with Momentum 实现
cpp
void SGD::doStep() {
for (size_t i = 0; i < parameters_.size(); i++) {
auto& param = parameters_[i];
auto grad = getDecayedGrad(param); // 加上权重衰减
if (!grad.defined()) {
continue;
}
if (momentum_ != 0.f) {
auto& buf = momentumBuffer_[i];
// 动量缓冲:buf = momentum * buf + (1-dampening) * grad
buf = buf.defined() ? (momentum_ * buf + (1.f - dampening_) * grad) : grad;
if (nesterov_) {
grad += momentum_ * buf; // Nesterov 动量
} else {
grad = buf;
}
}
*param -= lr_ * grad; // 参数更新
}
}
Adam 优化器实现
Adam 是深度学习中最常用的优化器之一,它结合了动量(Momentum)和 RMSprop 的思想:
cpp
void Adam::doStep() {
// 计算偏差校正因子
float b1t = 1.f - std::pow(beta1_, static_cast<float>(step_));
float b2t = 1.f - std::pow(beta2_, static_cast<float>(step_));
for (size_t i = 0; i < parameters_.size(); i++) {
auto& param = parameters_[i];
auto grad = getDecayedGrad(param);
if (!grad.defined()) continue;
auto& m = expAvg_[i]; // 一阶矩估计(动量)
auto& v = expAvgSq_[i]; // 二阶矩估计(RMSprop)
// 更新一阶矩
m = beta1_ * m + (1.f - beta1_) * grad;
// 更新二阶矩
v = beta2_ * v + (1.f - beta2_) * grad * grad;
// 偏差校正
auto mh = m / b1t;
auto vh = v / b2t;
// AdamW 权重衰减(解耦)
if (amsGrad_) {
auto& vMax = maxExpAvgSq_[i];
vMax = Tensor::max(vMax, vh);
*param -= lr_ * mh / (vMax.sqrt() + eps_);
} else {
*param -= lr_ * mh / (vh.sqrt() + eps_);
}
}
}
框架还实现了 Adagrad、RMSprop、Adadelta、AdamW 等优化器,它们都继承自同一个 Optimizer 基类,通过重写 doStep() 方法实现各自的更新规则。
分布式训练:NCCL 集合通信
对于大规模训练,单卡不够用时需要多卡甚至多机训练。框架通过 NCCL(NVIDIA Collective Communication Library)实现了分布式数据并行(DDP)训练。
数据并行原理
+----------+ +----------+
| GPU 0 | | GPU 1 |
| Model | | Model |
| (Rank 0) | | (Rank 1) |
+-----+----+ +-----+----+
| NCCL AllReduce |
+-------+--------+
|
梯度平均同步
数据并行的核心思路很简单:
- 每块 GPU 持有模型的完整副本
- 训练数据被切分到各 GPU(每块卡处理不同的 batch)
- 各 GPU 独立做前向传播和反向传播
- 通过 NCCL 的 AllReduce 操作,把所有 GPU 的梯度取平均
- 用平均后的梯度更新各 GPU 上的模型参数
NCCL 集成点
框架在 Optimizer 的 step() 方法之后、参数更新之前插入 AllReduce 通信:
cpp
for (auto& param : parameters) {
// 1. 各GPU本地计算梯度(已在backward中完成)
// 2. NCCL AllReduce:梯度平均
ncclAllReduce(param.grad().data(), param.grad().data(),
param.numel(), ncclFloat, ncclAvg, comm);
// 3. 本地参数更新
*param -= lr * param.grad();
}
运行分布式训练示例:
cpp
// 在两台机器上分别运行(rank 0 和 rank 1)
// Machine 1:
cd /examples/ddp/bin && ./example_ddp 0 0 2
// Machine 2:
cd /examples/ddp/bin && ./example_ddp 0 1 2
环境配置与编译运行
环境要求
- CMake:3.10 及以上版本
- C++ 编译器:支持 C++17 的编译器(g++ 7+、clang++ 5+、MSVC 2017+)
- CUDA Toolkit:11.0 及以上(可选,用于 GPU 支持)
- NCCL:可选,用于分布式训练
构建步骤
cpp
// 1. 创建构建目录
mkdir build
cd build
// 2. 配置 CMake(Release 模式,启用 CUDA 和 NCCL)
cmake -B . -DCMAKE_BUILD_TYPE=Release \
-DPROJECT_USE_CUDA=ON \
-DPROJECT_USE_NCCL=ON \
-DPROJECT_BUILD_EXAMPLES=ON \
-DPROJECT_BUILD_TEST=ON \
..
// 3. 编译
cmake --build . --config Release -j6
// 4. 运行测试
ctest --output-on-failure
CMake 构建选项
| 选项 | 默认值 | 说明 |
|---|---|---|
| PROJECT_BUILD_EXAMPLES | ON | 构建示例程序 |
| PROJECT_BUILD_TEST | OFF | 构建单元测试 |
| PROJECT_USE_CUDA | ON | 启用 CUDA 支持 |
| PROJECT_USE_NCCL | ON | 启用 NCCL 分布式训练 |
运行示例程序
cpp
// 自动微分基础示例
cd examples/autograd/bin && ./example_autograd
// nn.Module 使用示例
cd examples/module/bin && ./example_module
// 优化器使用示例
cd examples/optimizer/bin && ./example_optimizer
// MNIST 手写数字训练
cd examples/mnist/bin && ./example_mnist
// NCCL 集合通信测试(需要多GPU)
cd examples/nccl/bin && ./example_nccl <local_rank> <rank> <world_size>
// 分布式数据并行训练(需要多GPU)
cd examples/ddp/bin && ./example_ddp <local_rank> <rank> <world_size>
If you need the complete source code, please add the WeChat number (c17865354792)
完整训练示例:MNIST 手写数字识别
框架的 examples/mnist/ 目录下有一个完整的 MNIST 训练管线示例,展示了从数据加载、模型定义、训练循环到评估的完整流程。
训练流程大致如下:
- 数据加载:使用框架的 DataLoader 加载 MNIST 数据集,支持随机打乱、批量切分
- 模型定义:定义一个包含卷积层和全连接层的前馈网络
- 优化器配置:选择 Adam 优化器,设置学习率和权重衰减
- 训练循环 :
- 前向传播:输入图像经过网络得到预测
- 损失计算:用 CrossEntropyLoss 计算损失
- 反向传播:调用 loss.backward() 自动计算梯度
- 参数更新:调用 optimizer.step() 更新参数
- 梯度清零:调用 optimizer.zeroGrad() 清零梯度
- 评估:在每个 epoch 结束后在验证集上计算准确率
典型的训练输出:
Epoch 1/10, Loss: 0.2847, Val Acc: 0.9234
Epoch 2/10, Loss: 0.1456, Val Acc: 0.9512
Epoch 3/10, Loss: 0.0987, Val Acc: 0.9645
...
Epoch 10/10, Loss: 0.0312, Val Acc: 0.9823
落地用途
这个框架虽然"轻量",但功能相当完整,有几个实际的落地场景:
1. 深度学习教学与学习
对于想深入理解 PyTorch 底层原理的开发者来说,阅读和修改这个框架的代码比看十篇博客文章都管用。你可以清楚地看到:一个 tensor + tensor 的表达式是如何被翻译成计算图节点、如何分发到 CPU/GPU 内核、反向传播时梯度是如何沿图传递的。
2. 嵌入式与边缘部署
在资源受限的设备上(如 ESP32、树莓派),PyTorch Mobile 可能太重了。这个框架纯 C++ 实现、无外部依赖,可以方便地嵌入到嵌入式项目中。配合框架的序列化功能,还可以实现模型的保存和加载。
3. 自定义算子开发
如果你想研究新的神经网络算子(比如新的注意力机制、新的归一化方法),在这个框架里实现和调试比在大框架里方便得多。代码量小、依赖少,改完就能跑实验。
4. 大语言模型推理
框架已经支持 LLaMA、Qwen、Mistral 等主流 LLM 的推理。配合 CUDA 加速和混合精度,可以在消费级 GPU 上跑中等规模的模型。对于想做 LLM 本地部署的开发者来说,这是一个很好的起点。
5. 性能调优研究
由于框架完全可控,你可以方便地做各种性能实验:对比不同优化器的收敛速度、测试不同 batch size 下的吞吐量、分析 GPU 利用率等。这对于学习高性能计算和 GPU 编程也很有帮助。
总结与展望
从零手写一个深度学习框架是一项庞大的工程,但这个项目的代码结构清晰、分层合理,是学习框架内部原理的优秀教材。通过这个项目,你可以深入理解:
- 张量内存管理:自定义分配器、缓存策略、设备抽象
- 自动微分原理:计算图构建、Function 节点、反向传播链式法则
- GPU 编程:CUDA kernel 编写、CPU/GPU 分发机制
- 深度学习优化器:SGD/Momentum/Adam/AdamW 的数学原理和代码实现
- 分布式训练:数据并行、NCCL 集合通信
未来还可以继续扩展的方向包括:更多神经网络层(RNN、Transformer)、更多优化器(LAMB、Adafactor)、INT8 量化、GPU 内存优化、以及更完善的分布式训练支持(模型并行、流水线并行)。
如果你正在学习深度学习框架的内部原理,或者需要一个轻量级的 C++ 深度学习库嵌入到项目中,这个项目值得仔细研究。毕竟,理解"轮子"是怎么造出来的,才能造出更好的轮子。