5分钟跑通 tinygrad：一个小到能读完的深度学习框架

5分钟跑通 tinygrad一个小到能读完的深度学习框架【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygrad第一次接触 tinygrad你可能会被它的名字骗到——以为它只是 PyTorch 的缩水玩具。其实它是个完整能用的深度学习框架自带带自动求导的张量库、能把算子融合成单个 kernel 的编译器、JIT 加速还有训练神经网络要用的 nn 层和优化器。和 PyTorch 最大的不同在于它的编译器源码短到你能真的读完。如果你正在找一个不用背文档、改起来不心虚的训练框架它值得一试。它比 PyTorch 到底小在哪一句话PyTorch 像盖好的房子tinygrad 像给你图纸让你自己砌。你平时写 PyTorch 训练循环代码几乎能原样搬过来——TensorAPI、backward()、optim、数据集加载这些都在。但有几个体感上的差别没有 nn.Module模型就是一个普通 Python 类用__call__代替forward。想要卷积就直接写x.conv2d(w, b)不用为它专门开一个类。懒执行你写a b时其实什么都没发生要等realize()才真正算。看着怪但好处是相邻算子能被融成一个 kernel。看得见的编译器算子怎么变成 GPU 代码全程可查。右边就是 tinygrad 的全部家当框架层直接下到 kernel没有 cuDNN、cuBLAS 那一大坨黑盒。不想配环境直接跑起来要几步三步全程没有 CUDA 工具链的麻烦git clone https://gitcode.com/GitHub_Trending/tiny/tinygrad cd tinygrad python3 -m pip install -e .装完先确认它默认用哪块卡这行命令会打印CUDA、CPU、AMD之类的结果from tinygrad import Device print(Device.DEFAULT)确认无误后直接跑官方那个手写数字例子几秒内就能到 98% 准确率完整代码不到 50 行python3 examples/beautiful_mnist.py想要更细的分步教程看 MNIST 手册 或 快速上手 就行。第一次训练这么慢正常吗正常而且这恰恰是 tinygrad 的设计。现象你在 MNIST 手册 里按步骤训练一个 step 大概 75ms比 PyTorch 慢得多。原因tinygrad 没在派发速度上花心思它把加速的活儿交给了一个装饰器TinyJit。怎么解决给训练步套上 JIT同一个例子直接降到 1.0ms官方实测快了 75 倍from tinygrad import TinyJit jit_step TinyJit(step) # 或写成 TinyJit 装饰器原理很好理解前两次正常运行并把 kernel 录下来之后只回放 kernel、跳过 Python 开销。注意它会连优化器一起 JIT所以第二次之后会影响 kernel 的普通 Python 值会被冻住——Tensor的随机数不受影响照常工作。它真能跑大模型吗能showcase 里给的都是实打实的例子不是 demo 级别的YOLOv8python3 examples/yolov8.py直接出目标检测框。️WhisperSMALL1 python3 examples/whisper.py做语音转写。Stable Diffusionpython3 examples/stable_diffusion.py跑文生图。LLaMApython3 examples/llama.py能聊天还支持--quantize把权重压到 int8 或 nf4 塞进显存。也就是说视觉、音频、生成式、大语言模型这几条线它都有现成脚本权重格式统一走 safetensors能直接读 HuggingFace 的模型。换块卡要改代码吗不用。tinygrad 已经支持一批后端同一个Tensor代码切卡就行OpenCL、CPU、METAL、CUDA、AMD、NV、QCOM、WEBGPU。想自己接一块新卡门槛比想象中低——它只需要你实现大概 25 个底层算子。换句话说硬件适配的活儿是被刻意压得很薄的这也是它小而可改的底气。怎么看到它内部在干什么调试是 tinygrad 最顺手的一块几个环境变量就能打开透视镜完整列表见 环境变量说明DEBUG2 python3 examples/beautiful_mnist.py # 每个kernel的耗时、显存、带宽 BEAM2 python3 examples/beautiful_mnist.py # 自动搜最快kernel实现DEBUG2打印每个 kernel 跑多久、吃多少显存。DEBUG4直接看生成出来的 kernel 源码配合懒执行看融合效果最直观。BEAM2让它对 kernel 做束搜索挑你机器上最快的实现结果会缓存下次不用重搜。️VIZ1把计算图画出来调试网络结构特别有用。上手前必须知道的几件事⚠️它还没到 1.0API 已经稳定一阵了但偶尔还会动生产前多跑跑 test/ 里的回归测试。⚠️懒执行别忘 realize忘了realize()就print()大概率拿到的是个还没算的东西numpy()和item()会自动触发。️多卡就两行把模型和 batch 分别.shard到各卡即可官方在 MNIST 手册里留了道你能找到这两行吗的题。改 bug 是真能改的它纯 Python 且短官方明确鼓励你直接修库里的 bug 并补测试。最直接的下一步把仓库拉下来先跑通beautiful_mnist.py再打开DEBUG4盯着它生成的 kernel 看一遍——看懂那一刻你就真的拥有这个框架了。【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygrad创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考