tinygrad 完全解析极简、可 hack 的深度学习框架项目介绍介于 PyTorch 与 micrograd 之间的“玩具”内核在深度学习框架的版图上PyTorch 占据了易用性高地TensorFlow 统治着生产部署而 JAX 则代表着函数式编程的未来。然而这些框架的代码库动辄数十万行内部实现极其复杂让想要深入理解深度学习底层原理的开发者望而却步。tinygrad的出现正是为了填补这一空白。它由知名黑客 George Hotzgeohot创建是一个端到端深度学习堆栈包含自动微分、张量库、JIT 编译器、优化器、神经网络层和数据加载器但全部代码只有约 4 万行核心部分甚至更少。它的设计哲学是“保持小巧易于阅读极度可 hack”让你能够像理解 micrograd 一样彻底掌握一个现代深度学习框架的工作机制。tinygrad 的名字本身就暗示了它的定位比 PyTorch 更小比 micrograd 更完整。它不仅是教学工具更是一个真正可用的训练框架——你可以用它训练 MNIST 分类器98% 准确率仅需 5 秒也可以在其上构建 GPT 等现代模型。更重要的是它支持 CPU、GPUCUDA/Metal/OpenCL、AMD、Qualcomm 等多种硬件并通过 JIT 编译生成高性能内核。核心特性极简、统一、可编译轻量级张量库与自动微分tinygrad 的核心是Tensor类它支持所有常见的张量操作点积、reshape、卷积等并内置自动微分。与 PyTorch 的 API 高度相似让你可以无缝迁移代码。from tinygrad import Tensor x Tensor.eye(3, requires_gradTrue) y Tensor([[2.0, 0, -2.0]], requires_gradTrue) z y.matmul(x).sum() z.backward() print(x.grad.tolist()) # dz/dx print(y.grad.tolist()) # dz/dy函数式 JIT 与懒执行tinygrad 采用了与 JAX 相似的函数式 JIT 编译。通过TinyJit装饰器可以将一个函数编译为高效的设备内核。更重要的是tinygrad 是懒执行的——计算图不会立即执行而是先构建 IR最后通过编译器生成内核并在需要时触发计算。这种设计使得内核融合成为可能例如矩阵乘法可以被编译成单个高性能 kernel而不是多个串行操作。from tinygrad import Tensor, TinyJit TinyJit def matmul(a, b): return a b a Tensor.rand(1024, 1024) b Tensor.rand(1024, 1024) c matmul(a, b) # 触发 JIT 编译并执行多后端编译器tinygrad 的核心编译器负责将计算图降级为底层 IR并通过调度器为不同硬件生成代码。目前支持的后端包括CPUCUDANVIDIA GPUMETALApple GPUOpenCLAMDQCOMQualcommWebGPU浏览器添加新后端非常容易——只需实现约 25 个低级操作如矩阵乘法、卷积、复制等就能将整个框架移植到新硬件上。完整训练栈除了张量和自动微分tinygrad 还提供了nn模块线性层、卷积、BatchNorm 等、优化器SGD、Adam和数据加载器让你可以编写完整的训练脚本。from tinygrad import Tensor, nn from tinygrad.nn.optim import Adam class LinearNet: def __init__(self): self.l1 Tensor.kaiming_uniform(784, 128) self.l2 Tensor.kaiming_uniform(128, 10) def __call__(self, x): return x.flatten(1).dot(self.l1).relu().dot(self.l2) model LinearNet() optim Adam([model.l1, model.l2], lr0.001) # 训练循环类似 PyTorch for epoch in range(10): optim.zero_grad() loss model(x).sparse_categorical_crossentropy(y).backward() optim.step() print(loss.item())快速开始5分钟上手安装推荐从源码安装以获得最新特性git clone https://github.com/tinygrad/tinygrad.git cd tinygrad pip install -e .或直接通过 pip 安装主分支pip install githttps://github.com/tinygrad/tinygrad.git测试运行# 查看默认设备 python -c from tinygrad import Device; print(Device.DEFAULT) # 运行 MNIST 训练示例会下载数据需联网 python examples/beautiful_mnist.py调试模式设置环境变量DEBUG3可以查看编译过程DEBUG4可以看到生成的设备代码DEBUG3 python3 -c from tinygrad import Tensor; N1024; a,bTensor.empty(N,N),Tensor.empty(N,N); (a.reshape(N,1,N) * b.T.reshape(1,N,N)).sum(axis2).realize()优势对比tinygrad 与其他框架对比维度tinygradPyTorchJAXTensorFlow代码规模~4 万行数百万行数十万行数百万行可读性⭐⭐⭐⭐⭐⭐⭐⭐⭐⭐⭐自动微分支持支持支持函数式支持JIT 编译函数级 TinyJitTorchScript默认XLA硬件支持多后端可扩展丰富丰富丰富训练能力完整完整完整完整易用性接近 PyTorch极高中等中等生产级部署实验性成熟中等成熟社区新兴庞大增长庞大核心目标教学 可 hack易用 性能函数式 编译生产 部署核心优势代码极简任何人都可以在几小时内通读核心代码理解深度学习框架的内部机制。可 hack 性想要添加新操作、新后端或优化编译器修改几处代码即可无需在数十万行代码中挣扎。统一设计从张量到编译器的全栈统一没有“Python 前端 C 后端”的割裂感。真正跨平台通过少量底层原语就能移植到任何硬件甚至支持 WebGPU浏览器中运行。应用场景谁应该使用 tinygrad学习者想深入理解深度学习框架的内部实现tinygrad 是最佳教材。你可以从tinygrad/tensor.py开始一步步看到自动微分、编译和硬件加速如何实现。研究人员需要快速验证新算子、新优化器或新训练策略tinygrad 让你无需在庞大代码库中绕圈可以直接在核心上实验。框架开发者想为特定硬件如自研 NPU开发深度学习支持tinygrad 的硬件抽象层极其简单可以快速移植。极客喜欢探索新技术对 George Hotz 的作品感兴趣tinygrad 充满了 hack 精神值得收藏和把玩。总结学习深度学习框架的最佳入口tinygrad 证明了深度学习框架不必是庞然大物。通过极简的设计和可 hack 的实现它让开发者能够真正理解现代深度学习栈的每个层次从张量运算到编译器优化。如果你正在寻找一条从零开始掌握深度学习框架的捷径tinygrad 就是最好的起点。它可能不会取代 PyTorch 在生产环境中的地位但它一定会培养出更多懂得框架底层原理的工程师。现在就克隆代码用 DEBUG3 运行一下看看矩阵乘法是如何被编译成 CUDA kernel 的——你可能会对深度学习框架的魔力有全新的认识。项目地址https://github.com/tinygrad/tinygrad文档https://docs.tinygrad.org