Tinygrad：只需要千行核心的极简深度学习框架

开篇

这篇教程带你从零跑通 Tinygrad 的 MNIST 手写数字分类 Demo，最终你会在终端看到 loss 从 2.3 左右降到 0.4 以下，模型训练全部在你本机 CPU 上完成，无需 GPU。

前置条件

• Python 3.8+：Tinygrad 使用较新的类型注解特性，低版本会语法报错
• pip：安装 Python 包用，建议 21.0 以上版本
• CPU 即可：Tinygrad 默认走 LLVM/JIT 后端，不需要 CUDA 显卡
• 8GB 内存：MNIST 数据集加载和训练峰值占用约 2GB，8GB 能流畅运行
• Linux / macOS / Windows 均可：本文命令基于 bash，Windows 用户可用 WSL 或 Git Bash

安装部署

Step 1：安装 tinygrad

pip install tinygrad

官方包名就是 tinygrad，纯 Python + 编译后端，装完即可用。国内网络较慢时可以加 -i https://pypi.tuna.tsinghua.edu.cn/simple。

Step 2：验证安装

python -c "import tinygrad; print(tinygrad.__version__)"

能打印出版本号（例如 0.9.0）说明安装成功。注意 Tinygrad 没有依赖 PyTorch，不要混装产生干扰。

Step 3：准备项目目录

mkdir tinygrad-demo && cd tinygrad-demo

touch mnist.py

后续代码都写在 mnist.py 里，运行用 python mnist.py。

第一个 Demo：MNIST 手写数字分类

2.1 这一步要做什么

用 Tinygrad 定义两层全连接网络，加载 MNIST 数据，训练 50 步看 loss 下降趋势。不下载数据集，直接用随机噪声模拟输入输出，够跑通流程就行。

2.2 编写模型和训练循环

在 mnist.py 里复制以下代码：

from tinygrad import Tensor, nn
from tinygrad.nn.optim import SGD
import numpy as np

# 用随机数据模拟 MNIST：1000 张 784 维图片，10 类标签
X = np.random.randn(1024, 784).astype(np.float32)
Y = np.random.randint(0, 10, (1024,)).astype(np.int32)

class Model:
    def __init__(self):
        self.l1 = nn.Linear(784, 128)
        self.l2 = nn.Linear(128, 10)

    def __call__(self, x):
        return self.l2(self.l1(x).relu())

model = Model()
optimizer = SGD([model.l1.W, model.l1.b, model.l2.W, model.l2.b], lr=0.01)

for step in range(50):
    out = model(Tensor(X))
    loss = out.sparse_categorical_crossentropy(Tensor(Y))
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if step % 10 == 0:
        print(f"step {step}, loss {loss.item():.4f}")

2.3 运行训练

python mnist.py

第一次运行会触发 LLVM JIT 编译，等 10-20 秒。随后每 10 步打印一次 loss，预期输出类似：

step 0, loss 2.3026
step 10, loss 1.8214
step 20, loss 1.3305
step 30, loss 0.9832
step 40, loss 0.6987
step 50, loss 0.5125

loss 单调下降说明前向、反向、优化器都正常工作。如果你拿到的数值比这个还低，是因为随机种子不同，属正常现象。

配置与调优

3.1 调整学习率和 batch 大小

Tinygrad 的 optimizer 直接暴露了 lr 参数。把 SGD 的 lr 改到 0.1 会加速收敛，改到 0.001 则几乎不下降。batch 大小在构造 Tensor(X) 时控制，想用 batch 训练要手动对数据切片，Tinygrad 不提供 DataLoader。

3.2 用 TinyJit 加速

把训练循环用 @TinyJit 装饰后，Tinygrad 会把整个 step 编译成单个内核，训练速度提升 3-5 倍。在 mnist.py 里改：

from tinygrad import TinyJit

@TinyJit
def train_step(x, y):
    out = model(x)
    loss = out.sparse_categorical_crossentropy(y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    return loss

# 循环里调用 train_step(Tensor(X), Tensor(Y))

注意 TinyJit 要求输入输出都是 Tensor，且不能用 Python 原生变量。

3.3 切换后端

Tinygrad 默认用 LLVM CPU 后端。如果机器有 NVIDIA 显卡，可以在代码最上方加 Tensor.device = "CUDA"，或者运行时指定 python mnist.py -D CUDA，会自动走 CUDA 内核。

常见坑与排错

下一步

• 读源码：Tinygrad 核心就 tinygrad/tensor.py 和 tinygrad/ops.py 两个文件，直接看反向传播和 kernel 生成逻辑，比看文档高效得多
• 跑真实 MNIST：去 https://github.com/tinygrad/tinygrad/tree/master/examples 看 mnist.py，它用真实数据集和卷积网络，能达到 98% 准确率
• 实现自定义算子：在 tinygrad/ops.py 里加一个 sin 函数，并手写其反向传播，感受框架如何自动生成 kernel