Tinygrad — 极简深度学习框架，适合学习编译与推理原理

tinygrad：一个介于 [PyTorch](https://github.com/pytorch/pytorch) 和 [karpathy/micrograd](https://github.com/karpathy/micrograd) 之间的深度学习框架。由 [tiny corp](https://tinygrad.org) 维护。

[主页](https://github.com/tinygrad/tinygrad) | [文档](https://docs.tinygrad.org/) | [Discord](https://discord.gg/ZjZadyC7PK)

[](https://github.com/tinygrad/tinygrad/stargazers) [](https://github.com/tinygrad/tinygrad/actions/workflows/test.yml) [](https://discord.gg/ZjZadyC7PK)

tinygrad 是一个端到端的深度学习框架栈：

• 支持自动微分的 张量库
• 能够融合和降低内核的 中间表示（IR）和编译器
• 即时编译（JIT）+ 图执行
• 用于实际训练的 nn / optim / datasets 模块

它受到 PyTorch（人体工程学）、JAX（函数式变换和基于 IR 的自动微分）以及 TVM（调度和代码生成）的启发，但有意保持轻量级和可 hack 的特性。

tinygrad 对比

PyTorch

• ✅ 相似点：急切的 Tensor API、自动微分、optim、基础数据集和层。
• ✅ 你可以编写熟悉的训练循环。
• 🔁 与 PyTorch 不同，整个编译器和 IR 都是可见且可 hack 的。

JAX

• ✅ 基于 IR 的原始操作自动微分（类似 JAXPR + XLA）。
• ✅ 函数级 JIT（TinyJit），可捕获并重放内核。
• 🔁 函数式变换较少（目前没有完整的 vmap/pmap），但代码更易读。

TVM

• ✅ 多级降低、调度以及对内核的 BEAM 搜索。
• ✅ 用于批处理的设备“图”。
• 🔁 tinygrad 还提供了 前端框架（张量、nn、optim），而不仅仅是编译器。

惰性计算

尝试一个矩阵乘法。看看它如何通过惰性计算的威力，融合成一个单一的内核。

DEBUG=3 python3 -c "from tinygrad import Tensor;
N = 1024; a, b = Tensor.empty(N, N), Tensor.empty(N, N);
(a.reshape(N, 1, N) * b.T.reshape(1, N, N)).sum(axis=2).realize()"

我们可以将 DEBUG 改为 4 来查看生成的代码。

神经网络

事实证明，神经网络所需的 90% 功能是一个不错的自动微分/张量库。再加上一个优化器、一个数据加载器和一些计算，你就拥有了所需的一切。

from tinygrad import Tensor, nn

class LinearNet:
  def __init__(self):
    self.l1 = Tensor.kaiming_uniform(784, 128)
    self.l2 = Tensor.kaiming_uniform(128, 10)
  def __call__(self, x:Tensor) -> Tensor:
    return x.flatten(1).dot(self.l1).relu().dot(self.l2)

model = LinearNet()
optim = nn.optim.Adam([model.l1, model.l2], lr=0.001)

x, y = Tensor.rand(4, 1, 28, 28), Tensor([2,4,3,7])  # 替换为真实的 MNIST 数据加载器

with Tensor.train():
  for i in range(10):
    optim.zero_grad()
    loss = model(x).sparse_categorical_crossentropy(y).backward()
    optim.step()
    print(i, loss.item())

完整版本参见 examples/beautiful_mnist.py，该版本可在约 5 秒内达到 98% 的准确率。

支持的加速器

tinygrad 已支持多种加速器，包括：

• [x] OpenCL
• [x] CPU
• [x] METAL
• [x] CUDA
• [x] AMD
• [x] NV
• [x] QCOM
• [x] WEBGPU

并且很容易添加更多！你选择的加速器只需要支持总共约 25 个低级操作。

要检查默认加速器，请运行：python3 -c "from tinygrad import Device; print(Device.DEFAULT)"

安装

当前推荐的安装方式是从源码安装。

从源码安装

git clone https://github.com/tinygrad/tinygrad.git
cd tinygrad
python3 -m pip install -e .

直接安装（master 分支）

python3 -m pip install git+https://github.com/tinygrad/tinygrad.git

文档

文档和快速入门指南可以在 文档网站 上找到，该网站由 docs/ 目录构建。

与 PyTorch 对比的快速示例

from tinygrad import Tensor

x = Tensor.eye(3, requires_grad=True)
y = Tensor([[2.0,0,-2.0]], requires_grad=True)
z = y.matmul(x).sum()
z.backward()

print(x.grad.tolist())  # dz/dx
print(y.grad.tolist())  # dz/dy

在 PyTorch 中实现相同的功能：

import torch

x = torch.eye(3, requires_grad=True)
y = torch.tensor([[2.0,0,-2.0]], requires_grad=True)
z = y.matmul(x).sum()
z.backward()

print(x.grad.tolist())  # dz/dx
print(y.grad.tolist())  # dz/dy

贡献指南

最近 tinygrad 受到了很多关注。遵循以下指南将有助于你的 PR 被接受。

首先，以下情况会导致你的 PR 被关闭，并附上指向本节的说明：

• 禁止代码高尔夫！ 虽然低代码行数是本项目的指导原则，但任何看起来像代码高尔夫的代码都会被关闭。真正的目标是降低复杂性并提高可读性，删除 \n 对此毫无帮助。
• 所有文档和空白字符的更改 除非你是知名贡献者，否则将被关闭。编写文档的人应该是那些最了解代码库的人。没有证明这一点的人不应该修改文档。空白字符的更改既无用，又有引入错误的风险。
• 任何声称是“性能提升”的更改 都必须进行基准测试。一般来说，目标是简单性，因此即使你的 PR 使性能略有提升，也必须考虑可维护性和可读性之间的权衡。
• 通常，核心 tinygrad/ 文件夹之外的代码 没有经过充分测试，因此除非当前代码已损坏，否则你不应该更改它。
• 如果你的 PR 看起来“复杂”、差异很大或增加了大量代码行，将不会被审查或合并。考虑将其拆分为多个更小的、各自有明显优势的 PR。我常见的一种模式是在添加新功能之前进行先决条件的重构。如果你能（干净地）重构到该功能只需 3 行更改，那就太好了，这样我们审查起来也很容易。

现在，我们欢迎的是：

• 错误修复（附带回归测试） 很棒！这个库尚未达到 1.0 版本，所以如果你发现一个错误，修复它，编写测试，并提交 PR，这是有价值的工作。
• 解决悬赏任务！ tinygrad 提供现金悬赏 用于库的某些改进。所有新代码都应该是高质量且经过充分测试的。
• 新功能。 但是，如果你要添加一个功能，请考虑代码行数的权衡。如果是 3 行代码，那么它需要满足的有用性门槛就比 30 或 300 行代码要低。所有功能都必须有回归测试。一般来说，在没有其他约束的情况下，你的功能 API 应该与 torch 或 numpy 匹配。
• 有明显优势的重构。 一般来说，如果你的重构没有明显优势，它将被关闭。但有些重构非常棒！请从深层次核心角度考虑可读性。空白字符更改或移动几个函数是无用的，但如果你意识到两个 100 行的函数实际上可以通过参数使用同一个 110 行的函数，同时还能提高可读性，这就是一个巨大的胜利。重构应通过 过程回放测试。
• 测试/模糊测试。 如果你能添加非脆弱的测试，我们欢迎。我们这里也有一些模糊测试器，通过改进它们可以发现大量错误。发现错误，甚至编写带有 @unittest.expectedFailure 的（应该通过的）失败测试也很好。这是我们取得进展的方式。
• 从核心 tinygrad/ 文件夹中移除死代码。 我们不关心 extra 中的代码，但从核心库中移除死代码很棒。让新读者阅读和困惑的内容更少。

运行测试

你应该使用 pre-commit install 安装预提交钩子。这将在每次提交时运行 linter、mypy 和一部分测试。

有关如何运行完整测试套件的更多示例，请参考 CI 工作流。

一些本地运行测试的示例：

python3 -m pip install -e '.[testing]'  # 安装测试所需的额外依赖
python3 test/backend/test_ops.py        # 仅运行操作测试
python3 -m pytest test/                 # 运行整个测试套件

过程回放测试

过程回放 将你的 PR 生成的内核与 master 分支进行比较。如果你的 PR 是重构或性能提升，且没有预期的行为改变，它应该在拉取请求标题中包含 [pr]。