尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

10 分钟上手 tinygrad:5 秒训完 MNIST 的完整指南

10 分钟上手 tinygrad:5 秒训完 MNIST 的完整指南 10 分钟上手 tinygrad5 秒训完 MNIST 的完整指南【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygrad上周凌晨一点我要调一个张量融合的 bug开发机没 GPUPyTorch 的 wheel 光下载就花了 20 分钟装完还只能对着黑盒猜。换tinygrad之后安装只要 3 行命令纯 CPU 跑通 1024×1024 的矩阵乘法DEBUG多加一位数字就直接看到编译器生成的融合 kernel——前后不到 10 分钟。这就是你今天要跟我做的事。tinygrad 是介于 PyTorch 和 micrograd 之间的一个完整深度学习栈带自动求导的张量库、做 kernel 融合与 lowering 的 IR 和编译器、JIT 图执行外加nn/ optim / datasets 这些训练必备件。它故意做得小而可读核心代码全在tinygrad/目录里从 Tensor 到硬件指令每一步人眼都能跟完。当前版本 0.14.0要求 Python ≥ 3.11零硬性依赖。先跑起来3 条命令装完官方推荐从源码装一共就这几行git clone https://gitcode.com/GitHub_Trending/tiny/tinygrad cd tinygrad python3 -m pip install -e .这就是全部安装步骤装完先确认它会把活跑在哪块设备上python3 -c from tinygrad import Device; print(Device.DEFAULT)GPU 机器会显示CUDA、AMD、METAL之类纯 CPU 机器显示CPU。 装完先跑一次python3 -m pytest test/backend/test_ops.py当冒烟测试需要测试依赖pip install -e .[testing]确认后端没问题再碰模型能省掉后面一大半排查时间。你的第一个训练两层网络认出手写数字装好之后第一件事是训 MNIST别跳过去。官方例子 examples/beautiful_mnist.py 是 2 层卷积 2 层全连接batch 大小BS512默认STEPS70步测试集准确率 98%普通 CPU 机器上整个流程约 5 秒。训练循环的核心长这样from tinygrad import Tensor, nn, Context class Net: 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 Net(); opt nn.optim.Adam([model.l1, model.l2], lr0.001) x, y Tensor.rand(4, 1, 28, 28), Tensor([2,4,3,7]) with Context(TRAINING1): loss model(x).sparse_categorical_crossentropy(y).backward() opt.step()这是一个 10 步循环的压缩版完整训练逻辑直接看例子文件约 50 行。和 PyTorch 有两处你会立刻注意到的差别模型就是普通类没有nn.Module、没有forward只有__call__数据也简单from tinygrad.nn.datasets import mnist一行拿回 60000 张 28×28 的训练图和标签。写完一个矩阵乘法看看编译器对你做了什么tinygrad 的所有张量操作都是惰性的ab只是在图上挂了个节点直到你调.realize()或.numpy()才真正执行。好处是编译器有空间把相邻操作融合成一个 kernel。README 里这个一行命令就是干这个的——跑一个 1024×1024 matmulDEBUG3 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()跑完你会看到这次 matmul 被融合成了单个 kernel 以及参数搜索过程。DEBUG是分级的DEBUG2出每个 kernel 的耗时和带宽DEBUG3出优化过程DEBUG4出生成的 kernel 代码DEBUG5/6出 UOps 中间表示DEBUG7一路到汇编。完整清单在 docs/env_vars.md。想更狠一点可以加BEAM50让 beam search 给 kernel 多试几组参数。算在哪DEV 一个变量切换后端tinygrad 支持 OpenCL、CPU、METAL、CUDA、AMD、NV、QCOM、WEBGPU 等后端每个后端就是 tinygrad/runtime/ 里的一个ops_*.py文件。切换不用改代码运行前导出DEV就行DEVCL走 OpenCLDEVNV:CUDA:sm_70指定 NV 设备和 CUDA 目标架构DEVCPU:LLVM:x86_64,znver2,avx2精确到 CPU 微架构。格式是设备:渲染器:架构后两段可省略省略就自动探测。拿 YOLOv8 看效果最直接examples/yolov8.py 跑起来就是下面的实时检测画面从手写数字到 Stable Diffusion例子直接跑examples/目录基本是即开即玩的模型合集python3 examples/stable_diffusion.py出图SMALL1 python3 examples/whisper.py跑语音需要 pyaudio 和 torchaudiopython3 examples/llama.py聊天。识别图片这种最轻量的直接喂一张本地路径python3 examples/efficientnet.py ./test/models/efficientnet/Chicken.jpg。权重格式统一是 safetensorstinygrad/nn/state.py 里的get_state_dict/safe_save/safe_load/load_state_dict四个函数覆盖存读两行存、两行读也方便把别的框架的 safetensors 权重直接搬过来。FAQ被问得最多的 5 个问题Q我打印出来的数不对 / 跟 PyTorch 对不上为啥九成是惰性求值操作没 realize 之前图上什么都没发生。先加.realize()或.numpy()再核对两边的dtype是否一致。Q怎么确认代码跑在哪块设备上python3 -c from tinygrad import Device; print(Device.DEFAULT)或者干脆DEBUG1跑启动时会列出在用的设备。Q我想看编译器到底生成了什么 kernelDEBUG4打印生成的 kernel 代码DEBUG7打印目标硬件的汇编。看第一个例子建议从DEBUG3开始信息密度最舒服。Q推理慢了两倍先查什么给纯函数加TinyJit装饰器让它捕获并回放编译好的 kernelGPU 利用率会明显上去注意 JIT 不支持输入形状每次都变的函数。然后试BEAM50让 kernel 自动搜参数。Q模型训完怎么保存get_state_dict(net)拿参数字典safe_save(state_dict, model.safetensors)落盘加载反过来safe_loadload_state_dict两行。下一步你可以做什么跑一次完整测试python3 -m pytest test/确认你的后端全绿跟完 docs/quickstart.md 的 MNIST 教程再回头逐行读 examples/beautiful_mnist.py把训练循环和TinyJit对上号挑一个examples/里的模型跑通建议从python3 examples/stable_diffusion.py开始打开DEBUG4跑一个小模型第一次认真读完编译器给你生成的 kernel想深入原理的话翻 docs/developer/ 目录从 IR 和 uop 编译流程读起【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygrad创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表