tinygrad 快速上手：10 行代码跑通你的第一个深度学习训练-编程实验室

tinygrad 快速上手：10 行代码跑通你的第一个深度学习训练

【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygrad

你写过 PyTorch，也啃过 micrograd，但想要一个既够用、又小到能通读源码的深度学习框架？tinygrad 就是卡在两者之间的答案：一个纯 Python 写的端到端深度学习栈，张量、自动求导、编译器、JIT 全在你能读完的体量里。

tinygrad 到底能替你扛下什么活

大白话：tinygrad 让你用几十行 Python 就能定义、训练、跑推理一个神经网络，而且你能一路追到它生成的底层 kernel。技术上，它不是只给你一个张量库，而是一整套栈——带自动求导的 Tensor、负责融合和下推 kernel 的 IR 与编译器、把重复计算打包回放的 JIT，外加能直接训模型的 nn、optim、datasets。所以从手写 MNIST 到 Stable Diffusion、LLaMA 这类真实模型，它都能接。

7 行代码看 tinygrad 的自动求导

下面这段代码创建两个张量、做一次矩阵乘法求和，再算出梯度，是体验 tinygrad 核心——惰性张量加自动求导——最短的路径：

from tinygrad import Tensor x = Tensor.eye(3).clone() y = Tensor([[2.0, 0, -2.0]]) z = y.matmul(x).sum() z.backward() # 触发自动求导 print(x.grad.tolist()) # dz/dx

z.backward()是关键：它沿着计算图把每个张量上的梯度算出来，填进.grad。你不需要手写求导公式，clone()让 x 成为一个真实 buffer，这样梯度才有地方存。

懒加载：为什么 tinygrad 先记账再干活

可以把惰性求值想成餐厅点菜：你每下一道 op，服务员只是把菜写进单子，后厨不动；等你喊「上菜」（realize()），后厨才把能一起做的菜合成一锅、一次端上来。

a = Tensor.rand(4, 4) b = Tensor.rand(4, 4) c = (a * b).relu().sum() # 此时什么都没算，只记在账本上 c.realize() # 真正执行：融合、下推、出结果

这么设计是为了性能：一串 op 先攒起来，编译器才能把它们融合成更少的 kernel、减少在 CPU 和设备之间来回搬数据，顺带把整个计算图交给自动求导去处理。

从玩具代码到真训练：跑通一个数字分类器

你要实现一个手写数字分类器，识别 MNIST 里的 0 到 9。先搭一个带卷积的模型：

from tinygrad import Tensor, nn, Context from tinygrad.nn.datasets import mnist class Model: def __init__(self): self.conv = nn.Conv2d(1, 8, 5) self.fc = nn.Linear(8*24*24, 10) def __call__(self, x): return self.fc(self.conv(x).relu().flatten(1))

注意这里没有nn.Module，也没有forward：模型就是个普通 class，用__call__表达前向。接上优化器和数据，开始训练：

model, opt = Model(), nn.optim.Adam(nn.state.get_parameters(model), lr=3e-4) Xtr, Ytr, Xte, Yte = mnist() with Context(TRAINING=1): for i in range(100): opt.zero_grad() model(Xtr[:64]).sparse_categorical_crossentropy(Ytr[:64]).backward() opt.step()

get_parameters(model)会递归把模型里所有可训练张量抓出来交给 Adam；训练完在测试集上验一把：

acc = (model(Xte[:1000]).argmax(1) == Yte[:1000]).mean() print(acc.item())

想跑得更快，把前向套上 JIT，它会记住算好的 kernel 直接回放：

from tinygrad import TinyJit @TinyJit def forward(x): return model(x).realize()

新手最容易踩的 3 个坑

打印 Tensor 没有值 / 结果莫名不对现象：print(t)或拿结果去比较，发现数值是空的或错的。原因：op 只是记账，没realize()就不会真正计算。解法：用.numpy()、.item()或.realize()触发执行，再加DEBUG=3看生成了哪些 kernel。

训练态和推理态没区分，准确率不涨现象：加了 BatchNorm 的模型训练效果差、评估时对不上。原因：忘了开训练上下文，两种状态用了同一套统计量。解法：训练循环用with Context(TRAINING=1)包住，退出的时候会自动还原。

照着 PyTorch 习惯写代码报错现象：写了nn.Module或model.forward(x)找不到。原因：tinygrad 没有 Module，模型是普通 class，前向叫__call__。解法：用nn.state.get_parameters(model)取参数，直接model(x)调用。