
tinygrad 入门指南端到端深度学习栈的安装、Lazy 执行与实战训练【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygrad导读本文以 tinygrad 仓库根目录的 README.md 为主体系统讲解这个介于 PyTorch 与 micrograd 之间的端到端深度学习栈它既是带自动求导的 Tensor 库又是一套可融合、可下发的 IR 与编译器还内置 JIT 图执行以及 nn / optim / datasets 训练组件。读完本文你将掌握 tinygrad 的源码安装方式、用DEBUG环境变量观察惰性求值Laziness如何把一次矩阵乘融合成单个内核、如何写出与 PyTorch 等价的自动求导代码、如何跑起一个能逼近 98% 准确率的 MNIST 训练脚本以及如何在本仓库中运行测试来验证你的改动。一、tinygrad 是什么tinygrad 是一个端到端的深度学习技术栈由以下四层能力构成带 autograd 的 Tensor 库提供与 PyTorch 相近的 eagerTensorAPIIR 与编译器负责把算子融合fuse并逐层降低lower为可执行内核JIT 图执行函数级 JITTinyJit捕获并回放内核设备侧支持批量执行nn / optim / datasets面向真实训练的神经网络层、优化器与数据集工具。它受三套系统的启发PyTorch 的易用性ergonomics、JAX 的函数式变换与基于 IR 的自动微分functional transforms and IR-based AD、TVM 的调度与代码生成scheduling and codegen但刻意保持tiny体积小与hackable可随意修改。从 tinygrad/init.py 可以看到整个库的对外入口只暴露了极少的核心对象Tensor、TinyJit、function、Variable、dtypes、GlobalCounters、fetch、Context、getenv与Device——这就是tiny在代码层面的直接体现。与主流框架的定位对比维度PyTorchJAXTVMtinygradTensor API 与 autograd✅ eager API、optim、基础数据集与层✅ 基于原语的 IR 自动微分类似 JAXPR XLA仅编译器不提供前端框架✅ 三者的并集JIT图模式需额外 torch.compile✅ 函数级 JITTinyJit捕获并回放内核—✅ 函数级TinyJit编译器可见性 编译器与 IR 相对黑盒 可读性一般✅ 多层 lowering、调度、BEAM 搜索✅ 整个编译器与 IR 完全可见、可 hack功能变换—✅ 但 tinygrad 尚无完整vmap/pmap— 变换较少但代码极易阅读前后端一体✅✅❌ 只做编译器✅ 同时提供前端框架tensors、nn、optim与编译器一句话概括 README 的定位PyTorch 让你熟悉地写训练循环JAX 给你 IR 级自动微分TVM 给你调度与代码生成——tinygrad 则把这些都装进一个又小又能改的包里且连编译器本身都摊开给你看。二、Laziness一次矩阵乘如何融合成单个内核tinygrad 的核心执行模型是惰性求值Laziness你写的张量运算并不会立即触发计算而是先构建一张计算图IR等到显式realize()或item()、numpy()等需要具体数值的操作时才真正下发内核。由于整个计算在发布前可见编译器得以把多个算子融合进单个内核。README 给出了一个最直观的演示一个 1024×1024 的矩阵乘用reshape展开成 (N, 1, N) 与 (1, N, N) 的逐元素乘法再沿 axis2 求和这正是 matmul 的外积-求和形式然后.realize()触发编译。以DEBUG3运行可以看到调度与融合过程DEBUG3 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(axis2).realize()把DEBUG改成4则进一步打印生成的代码内核源码。DEBUG本身是一个在 tinygrad/helpers.py 中定义的ContextVar(DEBUG, 0)它从同名环境变量读取级别越高输出越详细DEBUG0默认静默DEBUG1打印打开设备、JIT 捕获等关键事件见 tinygrad/device.py 中opened device ...与 tinygrad/engine/jit.py 中JIT captured N linearsDEBUG2适合做计时统计examples/beautiful_mnist.py 中每轮训练前调用GlobalCounters.reset()正是为了配合DEBUG2观察每步耗时DEBUG3展示调度/融合信息DEBUG4展示生成的底层代码。类似的执行开关还有ContextVar家族例如BEAM内核 BEAM 搜索、NOOPT禁用优化、JIT是否启用 JIT等它们都定义在 tinygrad/helpers.py可以通过环境变量或Context上下文管理器在运行时切换。Context环境变量切换的运行时版本除了命令行环境变量tinygrad 还提供Context上下文管理器做运行时切换。以 README 训练示例中的TRAINING1为例它对应 tinygrad/helpers.py 中的TRAINING ContextVar(TRAINING, 0)。其实现tinygrad/helpers.py会在进入with块时保存旧值、写入新值退出时恢复from tinygrad import Tensor, nn, Context with Context(TRAINING1): for i in range(10): optim.zero_grad() loss model(x).sparse_categorical_crossentropy(y).backward() optim.step() print(i, loss.item())这种环境变量 / ContextVar 二选一的设计让超参数既能从 shell 注入适合批量实验也能在代码内安全地局部覆盖适合函数内切换训练/推理语义例如批量大小BS、步数STEPS、优化器选择SGD/MUON等均通过getenv读取见 examples/beautiful_mnist.py。三、用 Tensor API 写神经网络README 指出事实证明90% 的神经网络工作只需要一个像样的 autograd/tensor 库再加一个优化器、一个数据加载器和一些算力。 它给出的最小示例只用两个 Kaiming 初始化的权重矩阵就拼出了一个两层线性网络from tinygrad import Tensor, nn, Context 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], lr0.001) x, y Tensor.rand(4, 1, 28, 28), Tensor([2,4,3,7]) # replace with real mnist dataloader with Context(TRAINING1): for i in range(10): optim.zero_grad() loss model(x).sparse_categorical_crossentropy(y).backward() optim.step() print(i, loss.item())要点拆解Tensor.kaiming_uniform/Tensor.rand/Tensor.empty等构造方法定义在 tinygrad/mixin/creation.py如empty、full等均支持device与dtype参数.flatten(1)、.dot()、.relu()属于 tinygrad/mixin/movement.py 与 tinygrad/mixin/elementwise.py 提供的张量算子sparse_categorical_crossentropy与backward()构成完整的损失 反向传播闭环nn.optim.Adam等优化器定义在 tinygrad/nn/optim.py同一文件还提供SGD、AdamW、Muon等且所有优化器都遵循zero_grad()/step()的 PyTorch 风格接口便于写出熟悉的训练循环。README 特别指出完整的 MNIST 训练版本在 examples/beautiful_mnist.py可在约 5 秒内达到 98% 的准确率该数字来自仓库 README 的原始声明实际结果取决于硬件、随机种子与超参。这个示例还展示了 tinygrad 更进一步的工程化用法nn.Conv2d、nn.BatchNorm、nn.Linear、Tensor.max_pool2d组合成的卷积网络examples/beautiful_mnist.pyfunction装饰器把__call__变成可跟踪的复合函数TinyJit把train_step/get_test_acc编译为 JIT 图examples/beautiful_mnist.pynn.datasets.mnist(fashion...)直接从 URL 下载并解压 MNIST / Fashion-MNIST 数据tinygrad/nn/datasets.pygetenv(BS, 512)、getenv(STEPS, 70)、getenv(SGD)、getenv(MUON)等让批量大小、步数与优化器都可通过环境变量覆盖方便做基准与消融实验。运行完整训练脚本python3 examples/beautiful_mnist.py # Adam默认 70 步 python3 examples/beautiful_mnist.py # 可用 BS512 STEPS70 SGD1 等覆盖默认值四、加速器支持约 25 个底层算子走天下tinygrad 已支持众多加速器后端每个后端对应 tinygrad/runtime/ 目录下的一个ops_*.py模块加速器运行时模块OpenCLtinygrad/runtime/ops_cl.pyCPUtinygrad/runtime/ops_cpu.pyMETALtinygrad/runtime/ops_metal.pyCUDAtinygrad/runtime/ops_cuda.pyAMDtinygrad/runtime/ops_amd.pyNVtinygrad/runtime/ops_nv.pyQCOMtinygrad/runtime/ops_qcom.pyWEBGPUtinygrad/runtime/ops_webgpu.pyREADME 强调再加新的加速器很容易——你的加速器总共只需要支持约 25 个底层算子。 这是 tinygrad tiny 哲学的又一体现编译器把高层算子逐步降级到数量极少的底层原语后端只需实现这组原语即可接入。设备选择机制位于 tinygrad/device.pyALL_DEVICES列出自举设备列表tinygrad/device.pyDevice.DEFAULT会按可用性自动选择第一个可用的后端get_available_devices也可通过环境变量DEV或Context(DEV...)显式指定tinygrad/device.pyDevice[ix]按名称懒加载对应ops_*.py模块中的设备类并缓存实例。查看当前默认加速器python3 -c from tinygrad import Device; print(Device.DEFAULT)五、安装推荐从源码安装README 明确表示当前推荐从源码安装 tinygradgit clone https://github.com/tinygrad/tinygrad.git cd tinygrad python3 -m pip install -e .若只想快速体验最新 master也可以直接安装 GitHub 源码python3 -m pip install githttps://github.com/tinygrad/tinygrad.git注以上命令为 README 原始写法若在 gitcode 镜像环境下使用可将 clone 地址替换为对应镜像地址。仓库根目录的 pyproject.toml 定义了项目元数据与依赖pip install -e .即基于它做可编辑安装。六、快速示例与 PyTorch 的自动求导对比README 提供了一个同构示例展示 tinygrad 与 PyTorch 在 API 层面几乎一一对应——这正是PyTorch 易用性的直接证明。tinygrad 版本from tinygrad import Tensor x Tensor.eye(3).clone() # clone to make it a buffer y Tensor([[2.0,0,-2.0]]) z y.matmul(x).sum() z.backward() print(x.grad.tolist()) # dz/dx print(y.grad.tolist()) # dz/dyPyTorch 版本import torch x torch.eye(3, requires_gradTrue) y torch.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逐行对照可见tinygrad 的Tensor.eye(3)对应torch.eye(3, requires_gradTrue)但需要显式.clone()使其成为真实 buffer在 tinygrad 中JIT 输入也要求必须是真实 buffer请使用.clone()见 tinygrad/engine/jit.py.matmul()、.sum()、.backward()的调用方式一致.grad属性与.tolist()的读取方式一致。注意tinygrad 的Tensor默认即支持自动求导无需requires_gradTrue梯度在backward()后被写入各张量的.grad。七、文档与更多资源文档与快速上手指南构建自仓库的 docs/ 目录并发布为 docs 网站本地可用serve_docs.sh启动文档服务。docs 目录覆盖了多个主题例如docs/tensor/creation.md 等张量算子文档creation / elementwise / index / movement / ops / propertiesdocs/dtypes.md 数据类型说明docs/env_vars.md 环境变量总览对应上文提到的DEBUG、DEV、BEAM等docs/nn.md、docs/quickstart.md 神经网络层与快速开始docs/developer/ 面向开发者的实现细节layout、runtime、speed、uop 等。examples 目录还提供了大量可直接运行的模型示例如 examples/beautiful_mnist.py、examples/gpt2.py、examples/llama.py、examples/stable_diffusion.py、examples/yolov8.py 等是学习各模块用法的活教材。八、参与贡献与运行测试README 对贡献者有非常明确的态度与门槛核心要点如下会被关闭的 PR任何形式的 code golf为了缩短行数而牺牲可读性文档与空白字符改动除非是资深贡献者——写文档的人应该是最懂代码库的人声称提速但没有基准数据的改动改动核心tinygrad/目录之外未经充分测试的代码大而复杂的 PR建议拆成可独立评审的小 PR。受欢迎的贡献带回归测试的 bug 修复解决悬赏bounty任务所有新代码应高质量且有测试新特性API 尽量对齐 torch / numpy必须有回归测试明确的赢家式重构refactor 需通过 process replay 测试测试与模糊测试fuzzer包括用unittest.expectedFailure标记应该通过但当前失败的测试移除核心tinygrad/目录中的死代码。运行测试python3 -m pip install -e .[testing] # 安装测试所需的额外依赖 python3 test/backend/test_ops.py # 只跑算子测试 python3 -m pytest test/ # 跑整个测试套件建议先安装 pre-commit 钩子pre-commit install这样每次提交都会自动运行 linter、mypy 与部分测试。完整的测试矩阵可参考 CI 工作流 .github/workflows/test.yml。对于 Agent 场景README 建议始终用-n12并行运行测试以提速。Process replay 测试如果你的 PR 是重构或提速且不改变预期行为应当包含[pr]标记。process replay 机制会把你的 PR 生成的内核与 master 分支对比详见 test/external/process_replay/README.md确保行为完全一致。九、总结回到 README 开头的那句话——tinygrad 是介于 PyTorch 与 karpathy/micrograd 之间的东西。它的实用价值在于用 PyTorch 级别的 API 写训练代码用 JAX 级别的 IR 思维理解自动微分用 TVM 级别的调度与代码生成思路把算子融合进单个内核并且这一切都以小、可读、可 hack的方式装在一个库中。无论你是想快速跑通一个 MNIST 基准examples/beautiful_mnist.py还是想研究编译器如何把高层算子降低为约 25 个底层原语tinygrad/runtime/都可以从本文的安装、调试与训练路径出发在仓库中直接动手验证。【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygrad创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考