ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

tinygrad 深度学习框架上手:看懂从张量到 GPU 的完整流水线,以及后端怎么扩展

tinygrad 深度学习框架上手:看懂从张量到 GPU 的完整流水线,以及后端怎么扩展 tinygrad 深度学习框架上手看懂从张量到 GPU 的完整流水线以及后端怎么扩展【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygradtinygrad 是一个刻意保持小巧的端到端深度学习框架它把 PyTorch 风格的张量 API、自动求导和一条能从张量一路编译到具体硬件的流水线打包在一起。对熟悉 PyTorch 的人来说它更像 micrograd 和 PyTorch 之间的中间地带——既够用又小到你能把每一层都读明白。这篇文章带新手走查它的两条主线一条是可拦截的张量流水线lazy 操作 → 调度 → 融合 → 编译 → JIT另一条是可扩展的多后端同一套上层代码跑在 CPU、NVIDIA、AMD、Metal、WebGPU 上。读完你会知道怎么把一个小模型跑起来、怎么拆开看中间发生了什么、以及想接一块新硬件该动哪里。tinygrad 是什么一套能自己从张量写到硬件的深度学习栈先说结论tinygrad 不是一个只能调用的黑盒库而是一条你可以逐层打开的链路。它解决的核心问题是我想看清并改写张量到硬件之间发生的每一步。它由四块组成带自动求导的张量库tinygrad/tensor.py负责前向/反向的 API一套 IR 加编译器tinygrad/codegen/、tinygrad/schedule/负责把一堆小操作融合、降级成 kernelJIT 与图执行tinygrad/engine/jit.py负责把多次调用捕获、打包、重放nn / optim / datasetstinygrad/nn/、extra/datasets负责真实的训练流程。它借鉴了 PyTorch 的易用性、JAX 的基于 IR 的自动求导 函数级 JIT、以及 TVM 的调度与代码生成但刻意不做全功能保留够小、可改这一点。这张对比图能帮你建立预期你写的训练循环会和 PyTorch 很像但不同之处在于编译器与 IR 全程可见、可 hack。 适合谁看想理解深度学习框架内部到底怎么把x.matmul(w)变成一块 GPU kernel的人或者想给自己的设备加一个后端的人。不需要你懂 JAX 或 TVM只需要知道张量和反向传播是什么。核心机制张量流水线里每一层都能被拦截和改写这一章是全文的重点。tinygrad 里可拦截这件事对应的是你的张量操作默认是惰性的lazy只有当它被落地realize时才真正执行而落地之前要经过一连串可观察、可改写的层。具体流程是这样的你写z y.matmul(x).sum()此时只是在内存里搭一张操作图硬件上什么都没发生调用z.realize()或z.numpy()时图被交给调度器拆成一批 kerneltinygrad/schedule/编译器把这些 kernel 融合、降级成目标代码tinygrad/codegen/、tinygrad/renderer/运行时代码把它交给具体设备执行tinygrad/runtime/ops_*.py。关键在于第 2 到第 4 步之间的每一层都是可被拦截的你能让它停下来看也能改写它。README 里给了一个最直观的演示——用DEBUG环境变量控制打开到哪一层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设成3你会看到尽管写法是分开的几步惰性机制把它们融合进了一个 kernel再把它改成4直接能看到最终生成的代码。这就是 tinygrad 里拦截并观察的最小动作——不需要读几百行源码一个环境变量就能让你停在流水线的某一层看个明白。最小可用示例用十几行代码跑通一个训练循环把上面的机制落到代码上。下面是一个能跑的线性分类器结构上你会觉得眼熟因为它就该和 PyTorch 长得像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): 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]) # 实际项目换成真实数据加载 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())这段代码做了一件事前向model(x)→ 算损失 →backward()求梯度 →optim.step()更新参数循环往复。几个值得注意的点模型不需要继承任何基类__call__里就是对张量做一串操作。函数式操作如.relu()是方法而不是类这让网络定义非常直接。with Context(TRAINING1)把训练态打开退出上下文时会自动还原——类似 PyTorch 的model.train()但走的是全局上下文而不是对象状态。optim.zero_grad()放在循环里每轮先清零再更新避免梯度累加。完整版见examples/beautiful_mnist.py在 MNIST 上大约 5 秒能到 98% 精度。它能跑真实模型这一点下面这张图可以直接验证。后端怎么扩展一个 ops 接口对应多块硬件可扩展在 tinygrad 里最实在的体现是同一套上层张量代码能落到不同的硬件上而且加新硬件的门槛很低。设备相关的代码都集中在tinygrad/runtime/下每个ops_*.py就是一个后端实现ops_cpu.py、ops_nv.py、ops_amd.py、ops_cuda.py、ops_metal.py、ops_cl.py、ops_webgpu.py、ops_qcom.py等。换句话说可拦截流水线的最后一步把代码交给设备执行是被插件化拆开的——上层不知道下面到底是 CPU 还是某块 GPU。默认跑哪个后端可以直接问python3 -c from tinygrad import Device; print(Device.DEFAULT)想强制指定用DEV环境变量例如DEVCPU不同后端的接口选择也走这个变量。README 里给了一个很有用的数字一个新加速设备只需要支持大约 25 个底层 op就能接入 tinygrad。这是它可扩展承诺的量化表达——你不用重写整个框架只补一块接口实现。⚙️ 一个实用提醒tinygrad 也提供Tensor.from_blob这类零拷贝互操作接口可以和 PyTorch/OpenCL 的显存直接共享数据。用外部内存指针时要保证指针在整个张量生命周期内有效否则会读到坏数据——这在docs/runtime.md里有专门说明。进阶组合用 TinyJit 和图执行把多次调用打包重放当模型反复执行同样形状的前向时tinygrad 可以用 JIT 把开销降下来。用法很简单——给前向包一个装饰器from tinygrad import TinyJit TinyJit def jit(x): return net(x).realize() out jit(batch) # 第一次调用会捕获之后重放它背后的逻辑在tinygrad/engine/jit.py第一次调用时捕获kernel 序列之后把输入参数化成槽位、重放整条流水线如果某批 kernel 落在支持图执行的设备上还会被合并成一次图调用来批量执行。这里有个和参考架构里性能监控对应的思路——tinygrad 提供了几个开关帮你量和看VIZ1可视化计算图配合tinygrad/viz/观察流水线到底长什么样BEAM在编译阶段对 kernel 做更激进的搜索优化JITBEAM可单独控制tinygrad.helpers里的Timing给一段代码计时对比JIT 前 / JIT 后的加速。注意 JIT 的边界它目前不支持输入尺寸变化也不支持 tinygrad 之外的操作所以适合形状固定的推理循环而不是随意切换 batch 形状的训练。调试与验证DEBUG、VIZ 和 process replay 怎么用tinygrad 的验证方式很工程化核心是把改动是否改变了行为变成可执行的事逐层打开DEBUG3/4看融合与生成代码是调试中间件式的动作——停在流水线的某一层观察看图VIZ1看计算图判断融合、重排是否符合预期process replaytest/external/process_replay/会把你的分支生成的 kernel 和 master 对比。如果你做了一个只改结构、不该改结果的重构或加速PR 标题里带上[pr]就能用这套机制确认行为没漂移。这套机制的价值在于你不需要感觉它变快了而是有工具证明它确实按预期变快了、且没改坏别的。上面这张图来自test/models/efficientnet/展示的是用 tinygrad 跑一个真实视觉模型做分类这类验证场景的典型产物——这也是test/目录里大量测试在做的事情。 常见疑问挑几个新手容易踩的为什么我print(t)打印出来的不是我预期的值因为操作是惰性的print/.numpy()之前张量还没落地。想让它真正执行用.realize()或.numpy()。为什么 JIT 没加速检查输入形状是否每次都一样、以及前向里是否混入了 tinygrad 管不到的操作比如某些 Python 侧的控制流。加新后端要写多少大约 25 个底层 op 起步参考tinygrad/runtime/ops_python.py纯 Python 参考实现最直观。测试怎么跑装好.[testing]依赖后python3 -m pytest test/跑全量python3 test/backend/test_ops.py只跑 op 测试。上手路径给你一条最短路径避免一上来就啃uop/先跑通examples/beautiful_mnist.py确认环境和默认设备都 OK用print(Device.DEFAULT)看一眼打开DEBUG3/VIZ1重跑一次小模型看见流水线在做什么把前向套上TinyJit用Timing对比前后耗时体会打包重放的收益想接硬件时读docs/runtime.md和tinygrad/runtime/ops_python.py照着约 25 个 op 的接口补一个自己的ops_*.py。相关入口都是项目内的真实路径快速上手看docs/quickstart.md设备与接口看docs/runtime.md全部环境变量看docs/env_vars.md更多模型实现看extra/models想理解 op 语义看test/里的用例。需要 clone 源码时仓库地址是https://gitcode.com/GitHub_Trending/tiny/tinygrad本地python3 -m pip install -e .即可。一句话收尾tinygrad 的价值不在于比谁快而在于它把张量 → 硬件这条链路做成了可逐层拦截、可组合、可验证的结构——你既能当普通框架用它训练也能把它当成一台透明的机器去拆。【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygrad创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表