tinygrad 快速上手指南3 步跑通深度学习训练比 PyTorch 更轻量【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygradtinygrad 是一个把训练、编译、JIT、推理全塞进一个轻量库的深度学习框架主打代码短到你能读完、后端随便换。这篇写给刚从 PyTorch 过来、或者想搞懂深度学习底层怎么跑的新手。5 分钟跑通 MNIST 训练三步不用装任何额外东西。第 1 步克隆并安装git clone https://gitcode.com/GitHub_Trending/tiny/tinygrad cd tinygrad python3 -m pip install -e .第 2 步确认默认设备python3 -c from tinygrad import Device; print(Device.DEFAULT)有 NVIDIA 显卡会打印CUDA/NV没有就是CPU。想强制指定后端加个DEVCPU或DEVNV就行。第 3 步跑官方 MNIST 示例python3 examples/beautiful_mnist.py做完你能看到什么训练日志一行行往下刷大概 5 秒内准确率冲到 98% 以上。这就是 tinygrad 的完整训练循环——定义模型、算 loss、backward()求梯度、optim.step()更新参数和你熟悉的 PyTorch 几乎一个模子。下面这张图是 tinygrad 和其他框架的层次对比看完你就明白它轻在哪它到底怎么工作一张记账本的类比tinygrad 最核心的一个词是惰性lazy。你在 PyTorch 里写a b它马上就算。在 tinygrad 里写a b什么都不会发生。它只是往一张记账本上记一笔这里要加一次。你接着写* 2、.relu()它继续记。直到你调用.realize()或.numpy()、.item()这类要拿结果的调用它才把整本账翻出来一次性算完。打个比方你点了一桌菜厨房不会每道菜单独起锅送一次而是把所有步骤排好最后一下锅全炒出来端走。tinygrad 省下的正是每道菜单送那部分的开销。代码长这样from tinygrad import Tensor a Tensor.randn(1024, 1024) b Tensor.randn(1024, 1024) c (a * b).relu() # 还没算只是记账 c.realize() # 这里才上 GPU且乘法和 relu 融成了一个 kernel记账还有个好处编译器能把相邻的操作融合成一个 kernel。a*b和.relu()本来要来回显存搬两次数据融合后一次搞定带宽直接省下来。这就是它小模型快、还能读源码的底气。关键特性算子融合一次 realize 顶 N 次 kernel对你最大的好处是省显存带宽。GPU 很多时候不是算力不够而是在从显存搬数据上磨洋工。tinygrad 把能连起来的操作并成一个 kernel来回搬数据的次数直接砍掉。怎么实现的靠前面说的惰性记账加上 IR中间表示层面的 fuse 规则在realize()时一次性生成融合后的代码。想看它到底生成了什么跑DEBUG4就能看到 kernel 源码。TinyJit一行装饰器训练步从 75ms 干到 1ms这是 tinygrad 最值的功能。默认状态下它不算快——在 T4 上MNIST 一步训练要 75ms。把训练函数包一层TinyJit同一件事变成1.0ms约 75 倍加速。from tinygrad import TinyJit jit_step TinyJit(step) # step 是你的训练函数原理很简单它把第一次跑时用到的 kernel 都录下来后面只重放这些 kernel把 Python 那层开销整个跳掉。注意它连优化器一起 JIT相当于换个数据重放一遍。代价是非 tinygrad 的 Python 值从第二次开始会被冻住输入形状也不能变。后端随便挑同一份代码跑 CPU、NVIDIA、AMD、甚至浏览器你写好的模型换个DEV变量就能在别的硬件上跑代码一行不改。目前支持 OpenCL、CPU、Metal、CUDA、AMD、NV、QCOM、WebGPU浏览器这些后端。为什么这么省事因为每个新后端只要实现约 25 个底层 op其余全靠编译器生成。所以你能在自己的奇怪硬件上也跑起来后端源码 每个文件都不长。源码短到你能读完PyTorch 出问题你多半只能等官方修tinygrad 出问题你自己打开文件就能改。核心库是纯 Python量级比主流框架小得多。想加个功能、想改个 bug门槛就是读得懂。这也是它常被拿来当学习深度学习框架原理的教材的原因。配置详解这几个环境变量最常用tinygrad 的行为基本靠环境变量控制。下面是新手最常碰的几个直接照着用变量常用值调大/调小会怎样DEBUG1–7打印详情的开关。2看每个 kernel 耗时3看优化4看生成的 kernel 代码7看到汇编。调试时最有用。DEVNV/AMD/CPU/CL…指定后端。没显卡用CPU也能训DEVNV:CUDA:sm_70还能细到架构。BEAM整数如2、4越大搜索越狠能找出更快的 kernel 实现但首次搜索慢。结果会缓存之后不用再搜。DEFAULT_FLOATHALF/FLOAT32改HALF省显存、提速精度略降。默认FLOAT32。FLOAT161图像/视觉模型专用显存直接减半。JIT0/1/20全关纯调试1默认开2开 JIT 但关图执行。VIZ1打开可视化把计算图画出来排查结构超直观。TRAINING1切训练模式dropout 这类操作才会真正生效。举例想在浏览器里跑、还要看 kernel 代码就DEVWEBGPU DEBUG4 python3 xxx.py。完整清单在 env_vars.md。落地场景从 MNIST 到 Stable Diffusion场景一教学和快速原型几秒钟出结果痛点以前想给学生或自己演示一个训练流程得先装好 PyTorch CUDA环境一装就是半小时还容易版本冲突。tinygrad 的解法pip install -e .就行纯 Python没显卡也能在 CPU 上把 MNIST 跑到 98%。5 秒出结果用来讲原理、做原型验证环境成本几乎为零。场景二浏览器里跑视觉模型、生成式 AI痛点很多模型要部署到服务器上才能跑前端只能干等 API。tinygrad 有 WebGPU 后端把同一份模型直接塞进浏览器。比如仓库里的 Stable Diffusion 示例能在网页端直接出图目标检测同理。下面的 YOLOv8 结果就是 tinygrad 推理出来的能框出人、球等多个目标还有 EfficientNet 图像分类喂一张鸡的图它直接告诉你这是什么性能与调优把 tinygrad 调到和 PyTorch 一个量级tinygrad 默认不追求快但下面几步做完速度就能到主流框架的量级。数字以官方文档 T4 上的 MNIST 训练步为例配置单步耗时说明什么都不加~75ms调试够用别拿它跑正式训练TinyJit~1.0ms75 倍加速生产训练标配BEAM2再降一截自动搜最快 kernel结果缓存之后免搜FLOAT161视觉模型显存减半图像/生成类模型推荐配合 HALF 用可操作的几条永远上 TinyJit。这是最大的收益没有之一。训练函数写成独立函数再包TinyJit注意输入形状固定。首次跑加BEAM。它替你搜最快实现并缓存相当于免费的自动调优。视觉模型开FLOAT16/DEFAULT_FLOATHALF显存直接砍半大一点的网络才放得下。卡在哪用DEBUG2看。它给出每个 kernel 的耗时和带宽哪个 kernel 慢一眼就知道。常见问题默认设备怎么查python3 -c from tinygrad import Device; print(Device.DEFAULT)。要强制指定就用DEVNV这种环境变量。从 PyTorch 迁移要改什么两个点一是没有nn.Module模型就是普通类用__call__代替forward二是风格更功能式比如直接写x.conv2d(w, b)无状态的操作不用专门建类。API 大体和 torch 对齐迁移成本不高。为什么感觉比 PyTorch 慢默认状态它不优化速度。加TinyJit再配BEAM才进入和主流框架比拼的量级。模型权重怎么存标准格式是 safetensors和 Hugging Face 通用。用tinygrad/nn/state.py里的safe_save/safe_load存读即可。能跑 LLaMA 这种大模型吗能。仓库examples/llama.py就是完整的聊天示例下载好权重放进去就能跑。生成、检测、语音这些都有对应示例。写在最后tinygrad 的价值就一句话用最短的代码给你一条能从写模型一路看到GPU 上跑的汇编的完整链路。下一步建议跑完 MNIST 后把examples/beautiful_mnist.py打开对照着这份指南读一遍——你会发现整个训练循环也就几十行这正是它想让你感受到的东西。【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygrad创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考