tinygrad 文档精读：从源码安装、Tensor 惰性执行到与 PyTorch 的四大差异

tinygrad 文档精读从源码安装、Tensor 惰性执行到与 PyTorch 的四大差异【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygrad本篇文章以官方文档 docs/index.md 为骨架系统梳理 tinygrad 的安装方式、核心用法Tensor、dtypes、nn 库、TinyJit、多 GPU 分片以及它与 PyTorch 在编程范式上的关键差异并结合仓库源码如 examples/beautiful_mnist.py、tinygrad/nn/init.py补充实现细节。读完本文你将掌握 tinygrad 的基本编程模型知道如何安装、如何写一个可训练的模型并能理解函数式、惰性、需要显式 JIT这套设计哲学背后的原因。安装 tinygrad推荐从源码安装tinygrad 目前尚未发布 1.0 版本但官方文档明确指出API 已经保持了相当长一段时间的稳定。也就是说你现在学会的用法短期内不会轻易被推翻。安装有两种方式通过 pip 直接安装pip install tinygrad推荐从源码安装官方明确鼓励这种方式git clone https://github.com/tinygrad/tinygrad.git cd tinygrad python3 -m pip install -e .-e参数表示可编辑安装editable install这样当你后续阅读或修改仓库源码例如 tinygrad/tensor.py 或 tinygrad/nn/optim.py时改动会立即反映到已安装的包中非常适合学习与二次开发。安装完成后官方建议的下一步是先完成 MNIST 教程如果你是张量库新手可以通过 tinygrad-tensor-puzzles 这类张量解谜练习来熟悉张量运算思维想深入理解 tinygrad 内部原理的读者可以阅读 开发者文档。tinygrad 用法总览官方文档对 tinygrad 的日常使用给出了高度浓缩的概括核心要点如下。主类Tensor你打交道最多的类是 Tensor。它的行为与 PyTorch 非常相似但整体风格更偏向函数式。tinygrad 支持多种数据类型包括FLOAT32、HALF、BFLOAT16、FLOAT64、各类整型与无符号类型等具体集合以 tinygrad/dtype.py 中定义的DTypes为准。张量的创建方式与 PyTorch 高度一致例如 docs/quickstart.md 中的用法from tinygrad import Tensor, dtypes import numpy as np # 从已有数据创建 t1 Tensor([1, 2, 3, 4, 5]) t2 Tensor(np.array([1, 2, 3, 4, 5])) # 工厂方法 full Tensor.full(shape(2, 3), fill_value5) zeros Tensor.zeros(2, 3) ones Tensor.ones(2, 3) eye Tensor.eye(3) arange Tensor.arange(start0, stop10, step1) # 随机张量 rand Tensor.rand(2, 3) # 均匀分布 [0,1) randn Tensor.randn(2, 3) # 标准正态分布 uniform Tensor.uniform(2, 3, low0, high10) # 指定 dtype t3 Tensor([1, 2, 3, 4, 5], dtypedtypes.int32)更完整的工厂方法清单参见 Tensor 创建文档。所有操作都是惰性的这是 tinygrad 最核心的一条规则所有操作都是 lazy 的在你 realize 之前什么都不会真正执行。所谓 realize就是调用.realize()或.numpy()来触发计算t4 Tensor([1, 2, 3, 4, 5]) t5 (t4 1) * 2 t6 (t5 * t4).relu().log_softmax() print(t6.numpy()) # 直到这里才真正执行惰性求值带来的直接好处是tinygrad 可以看到整张计算图从而在 realize 时做算子融合、内存规划与调度优化这也是后面TinyJit能大幅提速的基础。完整的算子列表参见 Tensor Ops 文档。内置神经网络库 nntinygrad 自带一套 神经网络库包含常用层、优化器与权重存取管理网络层BatchNorm、Conv1d/Conv2d、ConvTranspose1d/ConvTranspose2d、Linear、GroupNorm、InstanceNorm、LayerNorm、LayerNorm2d、RMSNorm、Embedding、LSTMCell优化器SGD、LARS、AdamW、Adam、LAMB加载/保存safe_load、safe_save、get_state_dict、get_parameters、load_state_dict、tar_extract、torch_load以及 GGUF 权重加载gguf_load。从源码 tinygrad/nn/init.py 可以看到这些层是真正的普通类而非继承体系例如Linear的__init__只是用Tensor.uniform初始化weight与bias前向就是一行x.linear(self.weight.transpose(), self.bias)Conv2d的权重初始化采用scale 1 / sqrt(in_channels * prod(kernel_size))的均匀分布。理解这一点有助于你后续自定义层——你完全不需要继承任何基类。用 TinyJit 让计算变快tinygrad 内置 JIT即时编译只需在纯函数上加上TinyJit装饰器即可首次调用会录制用到的内核后续调用直接重放内核跳过 Python 层调度开销。官方文档特别强调tinygrad 不做 PyTorch 那样的快速 dispatch它把快这件事交给了 TinyJit因此要让程序跑得快必须显式使用它。多 GPU 支持tinygrad 对多 GPU 支持相当出色通过Tensor.shard即可把张量切分到多张卡上。训练时通常需要两件事把模型shard到所有 GPU以及按 batch 维度shard数据集。tinygrad Stack一张图看懂技术栈官方文档用下面这张图概括了 tinygrad 的技术栈全貌前端 Python API → 调度/图优化 → 各类后端这张图来自 docs/index.md可以帮助你理解 tinygrad 在表达能力与后端覆盖之间的取舍它用一套统一的张量抽象桥接 CPU、GPU 以及各类硬件后端。训练长什么样直接读 beautiful_mnist.py官方文档建议想理解 tinygrad 中的训练流程直接阅读 examples/beautiful_mnist.py。这个文件本身就是一份浓缩的官方示例其中包含几个值得注意的写法class Model: def __init__(self): self.layers: list[Callable[[Tensor], Tensor]] [ nn.Conv2d(1, 32, 5), Tensor.relu, nn.Conv2d(32, 32, 5), Tensor.relu, nn.BatchNorm(32), Tensor.max_pool2d, nn.Conv2d(32, 64, 3), Tensor.relu, nn.Conv2d(64, 64, 3), Tensor.relu, nn.BatchNorm(64), Tensor.max_pool2d, lambda x: x.flatten(1), nn.Linear(576, 10)] function def __call__(self, x:Tensor) - Tensor: return x.sequential(self.layers) TinyJit Context(TRAINING1) def train_step(self, X_train:Tensor, Y_train:Tensor) - Tensor: opt.zero_grad() samples Tensor.randint(getenv(BS, 512), highX_train.shape[0]) loss self(X_train[samples]).sparse_categorical_crossentropy(Y_train[samples]).backward() return loss.realize(*opt.schedule_step()) TinyJit def get_test_acc(self, X_test:Tensor, Y_test:Tensor) - Tensor: return (self(X_test).argmax(axis1) Y_test).mean()*100从中可以提取 tinygrad 训练的完整范式用普通列表组织前向x.sequential(self.layers)按顺序执行层列表无状态操作如Tensor.relu直接引用方法有状态层Conv2d、BatchNorm、Linear则是对象训练步骤写成函数并 jitTinyJit装饰训练步Context(TRAINING1)开启训练模式退出时自动恢复原值随机采样 batch用Tensor.randint在数据集内取索引配合X_train[samples]做索引采样一步到位loss.backward()计算梯度opt.schedule_step()返回需要 realize 的更新内核loss.realize(*...)一次性触发执行通过环境变量控制超参getenv(BS, 512)、getenv(STEPS, 70)使脚本可用BS256 STEPS100 python3 examples/beautiful_mnist.py这样的方式灵活配置。主循环则借助trange显示进度GlobalCounters.reset()让DEBUG2下的计时输出更干净训练中每隔 10 步在测试集上评估一次准确率。与 PyTorch 的差异四个关键点如果你是从 PyTorch 迁移过来的欢迎。官方文档说大部分 API 是相同的但你会发现 tinygrad 既熟悉又更接近正确的感觉。差异主要体现在四个方面。差异一tinygrad 没有 nn.Module在 tinygrad 中所谓Module没有任何特殊之处它就是一个普通类。PyTorch 靠继承nn.Module来递归收集参数而 tinygrad 用nn.state.get_parameters递归地扫描普通类里的合法 Tensor 即可。前向传播也不叫forwardtinygrad 直接使用__call__。class TinyNet: def __init__(self): self.l1 Linear(784, 128, biasFalse) self.l2 Linear(128, 10, biasFalse) def __call__(self, x): x self.l1(x).leaky_relu() return self.l2(x)对应地优化器直接接收参数列表SGD([net.l1.weight, net.l2.weight], lr3e-4)更省事的写法是SGD(get_parameters(net), lr3e-4)。差异二tinygrad 是函数式的在 tinygrad 里你可以直接调用x.conv2d(w, b)见 Tensor Ops 文档或x.sparse_categorical_crossentropy(y)所有无状态操作都没有对应的类。当然如果你希望有个地方存放状态tinygrad 也提供像 PyTorch 那样的Conv2d类。也就是说有状态的操作才需要类无状态的操作一律是函数。这一设计让层与层之间的组合极度自由例如beautiful_mnist.py中把Tensor.relu、Tensor.max_pool2d直接当作可调用对象塞进列表。差异三tinygrad 是惰性的a b在 tinygrad 里什么都不做直到你调用realize才真正执行计算。对比之下PyTorch 是 eager 模式每条语句立即执行。惰性让 tinygrad 能在执行前看到完整计算图从而做出全局优化代价是你必须时刻记得最后要 realize。差异四tinygrad 要求 TinyJit 才能快PyTorch 花了大量开发精力让 dispatch算子分派本身足够快tinygrad 不这样做。tinygrad 提供的是一个简单装饰器TinyJit会重放被装饰函数中用到的内核从而绕开 Python 层的调度开销。官方在 MNIST 教程 中给出过一个直观数据训练一步从约 75 ms 降到约 1 ms未同步 GPU第一个两次运行正常执行并录制内核从第三次运行起只重放 tinygrad 操作。同时要留意任何会影响内核的非 tinygrad Python 值会在第二次运行后被冻结所以 JIT 适用于输入输出是已 realize 张量、且不含动态分支的纯函数随机数生成类操作如Tensor.randint不受影响。补充控制运行行为的环境变量官方文档 docs/env_vars.md 列出了控制 tinygrad 运行行为的环境变量。它们既可以用在命令行如DEVCL DEBUG4 python3 -m pytest也可以用Context(...)装饰函数或用with Context(...)限定作用域例如from tinygrad import Context with Context(DEBUG0): a Tensor.ones(10, 10) a * 2核心全局变量速查变量取值说明DEBUG1–7调试输出见下方分级DEVAMD、NV、CPU、CL 等启用特定后端支持冒号分隔的设备:渲染器:架构三元组BEAM任意整数内核 beam search 的 beam 数量搜索多种实现取最快者并缓存DEFAULT_FLOATHALF、BFLOAT16 等指定默认浮点 dtype默认FLOAT32IMAGE1启用 2D 专用优化JIT0–20禁用1启用默认2启用但禁用图VIZ1启用计算图可视化ALLOW_TF321在 Ampere 及更新的 GPU 上启用 TensorFloat-32 张量核心CUDA_PATH路径字符串指定 CUDA 头文件目录其中DEV语法最特殊格式为设备:渲染器:架构各部分均可省略省略则由 tinygrad 自动推断还可用指定访问接口DEV取值含义AMD使用 AMD 设备AMD:LLVM使用 AMD 设备 LLVM 渲染器NV:CUDA:sm_70使用 NV 设备 CUDA 渲染器目标架构 sm_70AMD::gfx950使用 AMD 设备目标架构 gfx950USBAMD通过 USB 接口使用 AMD 设备CPU:LLVM使用 CPU 设备 LLVM 渲染器CPU:LLVM:x86_64,znver2,avx2,-avx512f同上并附带架构标志参见 runtime 文档DEBUG的分级含义DEBUG输出内容≥ 1启用调试并列出使用的设备≥ 2每个内核执行的性能指标耗时、内存、带宽≥ 3内核级别的优化应用情况≥ 4生成的内核代码≥ 5计算过程的中间表示 UOps≥ 6线性化后的 UOps 序列≥ 7面向目标硬件生成的汇编代码从文档到实战的路径最后把官方给出的学习路径汇总如下按顺序走完即可从会用进阶到能改快速入门docs/quickstart.md 假设你不懂 PyTorch手把手带你从建 Tensor 走到训练一个手写数字分类器含 JIT 加速、safetensors 保存/加载权重、VIZ1可视化计算图MNIST 教程docs/mnist.md 用 Keras 风格的卷积模型走完建模 → 取数 → 训练 → 提速全流程并展示了nn.optim.Adam、nn.state.get_parameters、Context(TRAINING1)与TinyJit的组合用法参考真实示例浏览 examples 目录如 examples/beautiful_mnist.py、examples/gpt2.py看完整模型如何组织深入源码tinygrad 是纯 Python 且代码量精简阅读 tinygrad 核心目录tinygrad/tensor.py、tinygrad/nn、tinygrad/device.py比修框架 bug 门槛低得多用例佐证test 目录下的测试是理解各算子语义与边界行为的第一手资料。记住贯穿全文的三句话Tensor 是核心、一切惰性、要快就 JIT。在此基础上函数式风格、无 nn.Module 的普通类、Tensor.shard多卡分片构成了 tinygrad 与 PyTorch 截然不同的编程体验。【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygrad创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考