一、先说说这玩意儿是干嘛的 学深度学习的人,十有八九是从PyTorch入门的。import torch,定义网络,调backward(),参数就更新了。方便是方便,但很多人用了三年PyTorch,却说不出autograd到底是怎么工作的——计算图长什么样?梯度是怎么从loss一层一层传回weight的?Tensor的内存是怎么管理的? 这个项目干的事,就是把这些黑箱全部打开。只用C++20标准库,写一个单头文件的张量库(约1400行),支持自动微分、矩阵乘法、激活函数、损失函数、优化器——全部自己实现。然后用这个库去训练GPT-2(124M参数),从数据加载到反向传播到参数更新,一条链路跑通。 最狠的是,你真的可以把这个头文件复制到你的项目里,include一下,就能开始写神经网络。没有CMake、没有Conda、没有pip install。 二、具体含义:它到底包含哪些东西? 模块 | 干啥的 | 为什么重要 | | tinytorch | 单头文件张量库(~1400行) | 核心,支持自动微分、各种算子 | | Tensor | 多维数组,带shape和stride | 所有运算的底座 | | Arena内存池 | 预分配一大块内存,Tensor从里面"切" | 零碎片、零malloc开销 | | 自动微分 | Forward记录计算图,Backward链式求导 | 不用手算梯度 | | MatMul | 矩阵乘法 | 模型的"肌肉" | | GELU | 激活函数 | 非线性变换 | | Softmax | 概率归一化 | 多分类输出 | | CrossEntropy | 交叉熵损失 | 训练目标 | | LayerNorm | 层归一化 | 训练稳定性 | | SGD/Adam | 优化器 | 参数更新 | | DataLoader | 加载TinyShakespeare数据集 | 训练数据 | | Tokenizer | BPE编码 | 文本Token ID | | GPT-2模型 | 12层Transformer | 完整的训练目标 |
三、代码实现原理:拆开来看 1. Arena内存池:所有Tensor从一个"大池子"里切 传统做法是每次创建Tensor都malloc一块内存。问题是:训练时创建销毁成千上万个中间Tensor,malloc/free的开销不可忽略,而且内存碎片严重。 这个项目的做法是:预分配一大块连续内存(Arena),所有Tensor从里面"切"出来。 class TensorContext { std::vector<std::byte> arena_; // 预分配的内存池 size_t arena_offset_ = 0; // 当前分配到的位置 public: TensorContext(size_t arena_size) : arena_(arena_size) {} Tensor* NewTensor(const std::vector<int>& shape, Dtype dtype = kFloat32) { // 从arena里切一块出来 size_t bytes_needed = NumElements(shape) * DtypeSize(dtype); std::byte* ptr = arena_.data() + arena_offset_; arena_offset_ += bytes_needed; // ... 初始化Tensor return tensor; }};创建Tensor时,arena_offset_往前跳,不需要malloc。训练循环结束后,整个Arena一次性重置,arena_offset_ = 0,所有中间Tensor"自动销毁"。 TensorContext ctx((size_t)1024 * 1024); // 1MB Arenaauto& W1 = *ctx.NewTensor({32, 3})->RandomNorm(); // 从Arena切一块auto& b1 = *ctx.NewTensor({32})->RandomNorm(); // 继续切// ... 训练结束后ctx.Reset(); // offset归零,全部回收没有指针释放的麻烦,没有内存泄漏的风险,而且所有Tensor在内存里是连续的,Cache友好。 2. Tensor:不只是数组,还是计算图的节点 这个项目的Tensor设计非常简洁,但功能完整: class Tensor { std::vector<int> shape_; // 形状 std::vector<int> stride_; // 步长 Dtype dtype_; // 数据类型 std::byte* data_ = nullptr; // 数据指针(指向Arena里的位置) std::byte* grad_ = nullptr; // 梯度指针(指向Arena里的位置) OpType op_ = OpType::kNone; // 这个Tensor是怎么来的 std::vector<tensor*> inputs_; // 父节点(谁生成了我)};每个Tensor节点记录: shape和stride:多维数组的索引方式 data和grad:值和梯度,都指向Arena里的内存 op:这个Tensor是怎么算出来的(MatMul?GELU?Softmax?) inputs:输入Tensor的指针(用于反向传播时找父节点) 3. 自动微分:Forward记录,Backward回溯 这是整个库的灵魂。你写的前向传播代码,其实同时在构建一张计算图。 // 构建计算图auto& hidden = (x.MatMul(W1) + b1).Gelu(); // hidden.op = kGelu, hidden.inputs = {add_result}auto& logits = hidden.MatMul(W2) + b2; // logits.op = kAdd, logits.inputs = {matmul_result, b2}auto& probs = logits.Softmax(); // probs.op = kSoftmax, probs.inputs = {logits}auto& loss = probs.CrossEntropy(y); // loss.op = kCrossEntropy, loss.inputs = {probs, y}每个操作返回一个新的Tensor,新Tensor记录了自己的op和inputs。整张图就是这样一层一层搭起来的。 Forward:从输入开始,沿着图往下算,直到loss。 void Tensor::Forward() { // 先算父节点 for (auto* input : inputs_) { if (input->op_ != OpType::kNone) { input->Forward(); // 递归 } } // 再算自己 switch (op_) { case OpType::kMatMul: MatMulForward(); break; case OpType::kGelu: GeluForward(); break; case OpType::kSoftmax: SoftmaxForward(); break; // ... }}Backward:从loss开始,沿着图往回传梯度。 void Tensor::Backward() { // 先算自己的梯度(根据op类型) switch (op_) { case OpType::kMatMul: MatMulBackward(); break; case OpType::kGelu: GeluBackward(); break; case OpType::kSoftmax: SoftmaxBackward(); break; // ... } // 再递归传回父节点 for (auto* input : inputs_) { input->Backward(); }}以MatMul为例,Forward和Backward的实现: // Forward: C = A @ Bvoid MatMulForward(Tensor* C, const Tensor* A, const Tensor* B) { int M = A->shape(0), K = A->shape(1), N = B->shape(1); float* c = (float*)C->data(); const float* a = (const float*)A->data(); const float* b = (const float*)B->data(); for (int i = 0; i < M; i++) { for (int j = 0; j < N; j++) { float sum = 0; for (int k = 0; k < K; k++) { sum += a[i * K + k] * b[k * N + j]; } c[i * N + j] = sum; } }}// Backward: dL/dA = dL/dC @ B^T, dL/dB = A^T @ dL/dCvoid MatMulBackward(Tensor* C, Tensor* A, Tensor* B) { float* grad_c = (float*)C->grad(); float* grad_a = (float*)A->grad(); float* grad_b = (float*)B->grad(); const float* a = (const float*)A->data(); const float* b = (const float*)B->data(); int M = A->shape(0), K = A->shape(1), N = B->shape(1); // dL/dA += dL/dC @ B^T for (int i = 0; i < M; i++) { for (int k = 0; k < K; k++) { for (int j = 0; j < N; j++) { grad_a[i * K + k] += grad_c[i * N + j] * b[k * N + j]; } } } // dL/dB += A^T @ dL/dC for (int k = 0; k < K; k++) { for (int j = 0; j < N; j++) { for (int i = 0; i < M; i++) { grad_b[k * N + j] += a[i * K + k] * grad_c[i * N + j]; } } }}注意+=而不是=:因为一个Tensor可能被多个下游节点使用(比如残差连接),梯度要累加。 4. GELU激活:不是简单的ReLU GPT-2用GELU而不是ReLU。GELU的公式是: GELU(x) = 0.5 * x * (1 + tanh(sqrt(2/π) * (x + 0.044715 * x^3)))比ReLU平滑,训练更稳定。Forward直接算,Backward需要链式法则: void GeluForward(Tensor* out, const Tensor* in) { float* o = (float*)out->data(); const float* x = (const float*)in->data(); for (int i = 0; i < in->NumElements(); i++) { float cdf = 0.5f * (1.0f + tanhf(0.7978845608f * (x + 0.044715f * x * x * x))); o = x * cdf; }}void GeluBackward(Tensor* out, Tensor* in) { float* grad_in = (float*)in->grad(); const float* grad_out = (const float*)out->grad(); const float* x = (const float*)in->data(); const float* o = (const float*)out->data(); for (int i = 0; i < in->NumElements(); i++) { float cdf = o / x; // 复用Forward结果 float pdf = expf(-0.5f * x * x) / sqrtf(2.0f * M_PI); float local_grad = cdf + x * pdf; grad_in += grad_out * local_grad; }}5. Softmax + CrossEntropy:合并优化 Softmax和CrossEntropy通常一起用。数学上,它们的组合可以简化: CrossEntropy(Softmax(logits), target) = -log(exp(logits[target]) / sum(exp(logits)))Backward时,Softmax的输出概率减去one-hot目标,就是梯度: void SoftmaxCrossEntropyBackward(Tensor* loss, Tensor* logits, int target) { float* grad_logits = (float*)logits->grad(); const float* probs = (const float*)logits->data(); // Forward时Softmax结果存在这里 for (int i = 0; i < logits->NumElements(); i++) { grad_logits = probs; // 先复制概率 } grad_logits[target] -= 1.0f; // 目标位置减1 // 这就是 dL/dlogits = probs - y_one_hot}这个技巧省掉了单独存Softmax中间结果的开销。 6. GPT-2模型:12层Transformer 用这个库搭GPT-2,代码读起来跟PyTorch几乎一样: // GPT-2 前向传播(简化)Tensor& GPT2Forward(Tensor& input_tokens, Tensor& target_tokens) { // Embedding auto& tok_emb = token_embedding.Embedding(input_tokens); // [B, T] -> [B, T, C] auto& pos_emb = position_embedding.Lookup(input_tokens); // 位置编码 auto& x = tok_emb + pos_emb; // 残差 // 12层 Transformer for (int layer = 0; layer < 12; layer++) { auto& ln1 = layers[layer].ln1.Forward(x); auto& attn = layers[layer].attn.Forward(ln1); auto& x1 = x + attn; // 残差连接 auto& ln2 = layers[layer].ln2.Forward(x1); auto& mlp = layers[layer].mlp.Forward(ln2); x = x1 + mlp; // 残差连接 } // 输出头 auto& ln_f = final_ln.Forward(x); auto& logits = lm_head.MatMul(ln_f); // [B, T, C] @ [C, V] -> [B, T, V] // 只取最后一个位置的logits算loss auto& last_logits = logits.Slice(-1); // [B, V] auto& probs = last_logits.Softmax(); auto& loss = probs.CrossEntropy(target_tokens); return loss;}每一层都是LayerNorm -> Attention -> 残差 -> LayerNorm -> MLP -> 残差的标准结构。 7. 训练循环:4步走 for (int step = 0; step < num_steps; step++) { // 1. 加载数据 dataloader.NextBatch(input_tokens, target_tokens); // 2. Forward auto& loss = GPT2Forward(input_tokens, target_tokens); loss.Forward(); // 3. Backward loss.ZeroGrad(); // 梯度清零 loss.Backward(); // 反向传播 // 4. 参数更新(Adam) for (auto* param : parameters) { float* w = (float*)param->data(); float* g = (float*)param->grad()->data(); float* m = (float*)param->m()->data(); // 一阶矩 float* v = (float*)param->v()->data(); // 二阶矩 for (int i = 0; i < param->NumElements(); i++) { m = beta1 * m + (1 - beta1) * g; // 动量 v = beta2 * v + (1 - beta2) * g * g; // 二阶矩 float m_hat = m / (1 - powf(beta1, step + 1)); // 偏差修正 float v_hat = v / (1 - powf(beta2, step + 1)); w -= learning_rate * m_hat / (sqrtf(v_hat) + eps); } } // 每40步验证一次 if (step % 40 == 0) { float val_loss = Evaluate(); printf("step %d, val loss: %f\n", step, val_loss); }}四、相关领域知识点 知识点 | 在这项目里的体现 | | 自动微分 | Forward建图,Backward链式求导 | | 计算图 | Tensor节点记录op和inputs,形成DAG | | Arena内存池 | 预分配,零碎片,Cache友好 | | 矩阵乘法 | 三重循环,OpenMP并行加速 | | GELU | GPT-2的标配激活函数 | | Softmax数值稳定性 | 减去最大值再exp,防止溢出 | | CrossEntropy梯度简化 | probs - one_hot ,省掉Softmax中间结果 | | LayerNorm | 减均值、除标准差、乘权重、加偏置 | | 残差连接 | x + f(x) ,防止梯度消失 | | Adam优化器 | 一阶矩、二阶矩、偏差修正 | | Causal Mask | Attention只算前面的token | | OpenMP | pragma omp parallel for 多核并行 |
五、设计思路:为什么这样设计? 设计选择 | 为什么这么做 | PyTorch会怎么做 | 单头文件 | 复制即用,零依赖 | 几百个文件,复杂的构建系统 | Arena内存池 | 零malloc开销,Cache友好 | 复杂的内存池和引用计数 | 自动微分 | 代码直观,跟PyTorch API类似 | 同样的autograd,但更复杂 | 递归Forward/Backward | 简单直接,容易理解 | 非递归的拓扑排序,更高效 | 标量循环矩阵乘法 | 可读性优先,OpenMP加速 | 调cuBLAS/MKL,黑盒 | FP32 | 训练精度优先 | 支持FP16/BF16混合精度 |
六、代码实现用途 学习自动微分的底层原理:比PyTorch的autograd简单100倍,但核心思想完全一致 理解GPT-2的完整训练链路:从数据加载到参数更新,每一行都是手写的 作为嵌入式项目的起点:单头文件,复制即用 教学演示:Arena、计算图、反向传播,全部可视化 性能基准:OpenMP加速后,CPU训练速度可观 七、流程原理图 图1:整体架构
If you need the complete source code, please add the WeChat number (c17865354792) 图2:tinytorch计算图与自动微分
图3:Arena内存池与训练循环
八、怎么跑起来?一步步来 1. 运行示例(不用下载数据) cd examplemake./example输出: step 0, val_loss: 2.97879step 1000, val_loss: 0.934616...step 29000, val_loss: 0.199498training loss: 0.00478648test case 0, x: 0.398134, 0.43387, 0.0558768, y_truth: 0, y_pred: 0 (correct)...这是一个简单的水果分类器(根据颜色、大小、重量判断是苹果、香蕉还是橙子),用来验证tinytorch能正常工作。 2. 下载GPT-2训练数据 ./dev/download_starter_pack.sh这会下载: GPT-2的预训练权重(HuggingFace格式) TinyShakespeare训练数据集 Tokenizer词汇表 总共约500MB。 3. 编译并训练GPT-2 make train_gpt2CPU训练(单线程): ./train_gpt2CPU训练(OpenMP多线程加速,推荐): OMP_NUM_THREADS=8 ./train_gpt2输出: [GPT-2]:max_seq_len: 1024vocab_size: 50257padded_vocab_size: 50304num_layers: 12num_heads: 12channels: 768Number of Parameters: 124475904...val loss: 5.32553step 0 train Loss: 4.67778 (took 7666.71 ms)step 1 train Loss: 5.19158 (took 7368.44 ms)...step 39 train Loss: 3.90174 (took 7558.28 ms)val loss: 4.29154generating:---Being barren savour, grantEveryone, every man and woman,Is heir to his life in unwritten words...---每步约7-8秒(取决于CPU)。训练40步后,模型就能生成看起来像莎士比亚的文本了。 4. 运行测试 张量操作测试(不需要数据文件): make test_tensor./test_tensorGPT-2前向/反向测试(需要下载数据): make test_gpt2./test_gpt2测试会验证: 矩阵乘法结果正确 GELU/Softmax/LayerNorm的Forward和Backward数值正确 GPT-2的loss跟参考实现一致 5. 在自己的项目里使用tinytorch ...int main() { // 创建Arena(1MB) tt::TensorContext ctx((size_t)1024 * 1024); int variables = 3, neurons = 32, classes = 3; // 输入和目标 auto &x = *ctx.NewTensor({1, variables}); auto &y = *ctx.NewTensor({1}, tt::kI32); // 参数 auto &W1 = *ctx.NewTensor({neurons, variables})->RandomNorm(); auto &b1 = *ctx.NewTensor({neurons})->RandomNorm(); auto &W2 = *ctx.NewTensor({classes, neurons})->RandomNorm(); auto &b2 = *ctx.NewTensor({classes})->RandomNorm(); std::vector<tt::tensor *=""> params = {&W1, &b1, &W2, &b2}; // 构建计算图 auto &hidden = (x.MatMul(W1) + b1).Gelu(); auto &logits = hidden.MatMul(W2) + b2; auto &probs = logits.Softmax(); auto &loss = probs.CrossEntropy(y); // 训练 float learning_rate = 0.001; for (int step = 0; step < 30000; step++) { // 填充数据 x.Fill({0.5, 0.3, 0.8}); // 颜色、大小、重量 y.Fill({1}); // 标签 loss.Forward(); loss.ZeroGrad(); loss.Backward(); // SGD更新 for (auto param : params) { auto weights = (float *)param->data(); auto grads = (float *)param->grad()->data(); for (int i = 0; i < param->NumElements(); i++) { weights -= learning_rate * grads; } } } // 推理 x.Fill({0.5, 0.3, 0.8}); probs.Forward(); auto result = probs.Flatten(); // 各类概率 return 0;}6. 常见问题 编译报错找不到OpenMP:去掉Makefile里的-fopenmp和-lomp,或者安装libomp 训练loss不下降:检查学习率,或者多跑几步 内存不够:减小batch size或seq_len 生成文本乱码:确认tokenizer和权重匹配 九、总结 这个项目最大的价值,是让你看到一个"能训练"的GPT-2可以有多小、多直白。 它用一个单头文件(~1400行)实现了: 完整的张量系统(多维数组、shape、stride) 自动微分引擎(Forward建图、Backward链式求导) 常用算子(MatMul、GELU、Softmax、CrossEntropy、LayerNorm) 优化器(SGD、Adam) GPT-2的完整训练链路(12层Transformer) 而且全部代码加起来可能就几千行,没有PyTorch的层层封装,没有CUDA的复杂配置,每个Forward()、每个Backward()、每个arena_offset_都清清楚楚。 如果你想真正理解"深度学习框架的底层是怎么工作的",而不是只会调包,这个项目就是最好的起点。把tinytorch通读一遍,亲手改改网络结构、学习率、激活函数,你会对自动微分和Transformer训练有一个完全不一样的体感。 Welcome to follow WeChat official account【程序猿编码】 </tt::tensor></tensor*> 免责声明:如果侵犯了您的权益,请联系站长,我们会及时删除侵权内容,谢谢合作! |