企业官网建设流程全解析

tinygrad MNIST 入门教程：从模型定义、训练到 TinyJit 提速的完整实战

【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygrad

MNIST 手写数字识别是深度学习框架的"Hello World"。本教程以 docs/mnist.md 为主线，带你用 tinygrad 从零完成一个可跑到 98%～99% 准确率的卷积网络：包括环境搭建、默认设备确认、模型定义、数据集加载、随机采样训练，以及 tinygrad 最具特色的TinyJit编译加速。读完你将掌握 tinygrad 的"纯 Python 类即模型 + 惰性求值 + 无 Python 开销的 JIT 回放"这套核心工作流，并能在 examples/beautiful_mnist.py 等仓库示例的基础上进一步自由发挥。

安装与环境准备

tinygrad 非常轻量，安装"随处可行"。最简单的方式是直接用 pip 从 Git 仓库安装（例如在本地 notebook 或 Google Colab 中）：

!pip install git+https://github.com/tinygrad/tinygrad.git

Colab 上不需要任何特殊的镜像，但为了速度建议选择 T4 GPU 类型的运行时。

安装完成后，可以用下面几行确认默认设备：

from tinygrad import Device print(Device.DEFAULT)

在 GPU 实例上你会看到CUDA，在纯 CPU 实例上则是CPU。从源码看，Device.DEFAULT的定义在 tinygrad/device.py 中，其取值逻辑为DEV.device or self._select_device——即当没有显式指定DEV环境变量时，框架会自动探测并选择当前可用的后端（CUDA、CPU、Metal、AMD 等）。tinygrad 支持通过DEV环境变量或with Context(DEV=...)来切换默认设备，且设置Device.DEFAULT本身已被标记为弃用。

一个简单的卷积模型

我们采用 Keras 官方 MNIST 卷积示例中的经典结构：两层 Conv2d 加一层 Linear，中间用 ReLU 和 MaxPool2d 降采样。

from tinygrad import Tensor, nn, Context class Model: def __init__(self): self.l1 = nn.Conv2d(1, 32, kernel_size=(3,3)) self.l2 = nn.Conv2d(32, 64, kernel_size=(3,3)) self.l3 = nn.Linear(1600, 10) def __call__(self, x:Tensor) -> Tensor: x = self.l1(x).relu().max_pool2d((2,2)) x = self.l2(x).relu().max_pool2d((2,2)) return self.l3(x.flatten(1).dropout(0.5))

与 PyTorch 相比有两个关键差异：

• 只有带状态的层（卷积、全连接）在__init__中声明；ReLU、MaxPool2d、Dropout 这类无参数操作直接用Tensor的方法链式调用；
• 没有nn.Module基类，也没有forward方法，就是一个普通 Python 类加__call__。

从仓库源码可以印证这一点：tinygrad/nn/init.py 中Conv2d.__init__只是创建self.weight与self.bias两个Tensor（用Tensor.uniform按1/sqrt(in_channels * prod(kernel_size))的尺度初始化），__call__直接调用x.conv2d(self.weight, self.bias, ...)；Linear同理，__call__是x.linear(self.weight.transpose(), self.bias)。所谓"层"本质上是持有参数的普通类，这正是nn.state.get_parameters能通过遍历对象属性自动收集参数的原因（见 tinygrad/nn/state.py 的get_state_dict/get_parameters，它递归遍历对象__dict__、list、tuple、dict，把叶子Tensor收集为参数字典）。

注：nn.Conv2d也支持字符串形式的 padding，如padding='same'，但要求stride=1，否则会抛出ValueError（见 tinygrad/nn/init.py）。

获取 MNIST 数据集

tinygrad 自带 MNIST 数据集加载，只需一行：

from tinygrad.nn.datasets import mnist X_train, Y_train, X_test, Y_test = mnist() print(X_train.shape, X_train.dtype, Y_train.shape, Y_train.dtype) # (60000, 1, 28, 28) dtypes.uchar (60000,) dtypes.uchar

返回四个张量：训练图像、训练标签、测试图像、测试标签。图像 shape 为(60000, 1, 28, 28)，标签 shape 为(60000,)，dtype 都是dtypes.uchar（无符号 8 位整数）。

实现非常简单（tinygrad/nn/datasets.py 一共只有几行）：mnist()通过Tensor.from_url直接下载官方 MNIST 的四个 gz 文件，跳过 16 字节的 IDX 文件头（[0x10:]）后reshape(-1,1,28,28)成图像张量，标签则跳过 8 字节头。它还支持fashion=True参数切换到 Fashion-MNIST 数据集，以及device=None参数指定加载到哪个设备。同一文件里的cifar()则展示了如何用tar_extract解包 CIFAR-10 的 tar 压缩包，可作为参考。

模型推理：惰性求值

由于 MNIST 足够小，mnist()会把数据集直接复制到默认设备上。于是"创建模型 + 跑一次推理 + 算准确率"只需要：

model = Model() acc = (model(X_test).argmax(axis=1) == Y_test).mean() # NOTE: tinygrad is lazy, and hasn't actually run anything by this point print(acc.item()) # ~10% accuracy, as expected from a random model

这里有一个 tinygrad 的核心心智模型：惰性求值（lazy evaluation）。model(X_test).argmax(...).mean()只是构建了一张计算图，此刻没有任何计算真正发生；直到调用.item()（或.realize()）才触发实际的 kernel 执行。随机初始化模型的准确率约 10%（10 类均衡分布的期望值），完全符合预期。

这一设计贯穿整个框架：例如 tinygrad/tensor.py 中的realize(*lst)负责真正触发计算，而backward()（tinygrad/tensor.py）同样只是构建反向计算图。在后面的训练中，我们会看到"惰性 + 显式 realize"如何与TinyJit配合，把 Python 层的调度开销完全消除。

训练模型：Adam + 随机采样

训练使用 Adam 优化器。nn.state.get_parameters(model)会遍历模型类、把所有参数（weight/bias）收集出来交给优化器。tinygrad 的习惯是把单步训练写成一个函数，以便后续用TinyJit编译：

optim = nn.optim.Adam(nn.state.get_parameters(model)) batch_size = 128 @Context(TRAINING=1) def step(): samples = Tensor.randint(batch_size, high=X_train.shape[0]) X, Y = X_train[samples], Y_train[samples] optim.zero_grad() loss = model(X).sparse_categorical_crossentropy(Y).backward() optim.step() return loss

逐行拆解：

• @Context(TRAINING=1)：把TRAINING这个全局上下文变量设为 1。从 tinygrad/helpers.py 可见，TRAINING = ContextVar("TRAINING", 0)，而 tinygrad/nn/optim.py 的Optimizer.schedule_step会检查if not TRAINING: raise RuntimeError(...)——不开启 TRAINING 直接调用优化器会直接报错，提示你使用Context(TRAINING=1)。这也让nn.BatchNorm等层能区分训练/推理统计量。
• Tensor.randint(batch_size, high=X_train.shape[0])：从 0 到 59999 之间随机采样 128 个样本索引（tinygrad/mixin/rand.py 中的randint类方法，low默认为 0）。
• X_train[samples]：用整数索引张量做花式索引取批次。
• sparse_categorical_crossentropy(Y)：稀疏交叉熵损失，Y是整数标签（tinygrad/mixin/op.py 中的实现，支持ignore_index、label_smoothing、reduction参数，默认reduction="mean"）。
• .backward()：自动微分，构建并执行反向传播。
• optim.zero_grad()/optim.step()：清空梯度、更新参数。

计时单步：

import timeit timeit.repeat(step, repeat=5, number=1) #[0.08268719699981375, # 0.07478952900009972, # 0.07714716600003158, # 0.07785399599970333, # 0.07605237000007037]

T4 Colab 上约 75 ms 一步。

如果想看每个 kernel 的时间分解，可以用：

from tinygrad import GlobalCounters, Context GlobalCounters.reset() with Context(DEBUG=2): step()

DEBUG是 tinygrad 最常用的调试开关之一（tinygrad/helpers.py 中DEBUG = ContextVar("DEBUG", 0)）：

• DEBUG=2显示每个被运行的 kernel 及其耗时；
• DEBUG=4会进一步打印生成的 kernel 源码；
• 所有的Context变量同时也都能通过同名环境变量设置（例如DEBUG=2 python train.py）。

关于优化器：仓库 tinygrad/nn/optim.py 按复杂度排序提供了完整家族——SGD（本质是关闭 trust ratio 的LARS）、Adam/AdamW（本质是关闭 trust ratio 的LAMB）、Muon（LARS + Newton-Schulz 迭代，ns_steps=5）等。Adam默认lr=0.001, b1=0.9, b2=0.999, eps=1e-8；优化器还会把学习率提升到至少 float32 精度存储。

为什么这么慢？——引入 TinyJit

75 ms 一步对调试来说足够快，但对训练来说并不理想。与 PyTorch 不同，tinygrad 天生不是为"这种慢速逐步执行"设计的。这里引入 tinygrad 第一个标志性概念：TinyJit。

from tinygrad import TinyJit jit_step = TinyJit(step)

注意它也可以用作装饰器：@TinyJit。再计时：

import timeit timeit.repeat(jit_step, repeat=5, number=1) # [0.2596786549997887, # 0.08989566299987928, # 0.0012115650001760514, # 0.001010227999813651, # 0.0012164899999334011]

从 75 ms 降到 1.0 ms，约 75 倍加速。注意这里没有显式同步 GPU，所以真实 GPU 耗时可能略高于这个数字。

TinyJit的工作原理（可从 tinygrad/engine/jit.py 的_TinyJit.__call__实现完整印证）：

1. 前两次调用正常执行并捕获：第一次调用（cnt==0）按原样运行；第二次（cnt==1）进入"capture"阶段，用capturing上下文把执行过程中触发的所有 kernel（Ops.LINEAR序列）记录下来，然后经jit_lower做参数化（输入 buffer 映射为 PARAM）、内存规划（memory_plan_rewrite）和可选的 graph 化（graph_split_rewrite，受JIT_BATCH_SIZE控制，默认 32），最终打包成一个CapturedJit。
2. 从第三次起只做"回放"：跳过整个 Python 执行过程，只按捕获的 kernel 序列在 GPU/CPU 上重放，Python 调度开销被完全消除。这就是"dumb replay on different data"——把 JIT 当作对不同数据的无脑重放。
3. 注意事项：任何影响 kernel 的非 tinygrad Python 值都会在第二次运行后"冻结"（因为回放不再经过 Python）。不过Tensor的随机函数（如Tensor.randint）是编译进 kernel 的，因此能正常工作——这也是上面step()里随机采样放在函数内部仍然有效的原因。
4. 与大多数框架不同，tinygrad 的 JIT连优化器也一起 JIT（optim.step()也在被捕获的函数内）。

需要理解Tensor输入的规范化约定：_prepare_jit_inputs要求 JIT 函数的 Tensor 参数必须是真实 buffer（虚拟视图会报JIT inputs must be real buffers; use .clone()），且不允许重复输入。若捕获期间生成了惰性.grad等，held_bufs收集逻辑会保证它们不被内存规划器误释放。另外_TinyJit还支持reset()（重新捕获）与prune=True（剪枝掉与输出无关的一次性 kernel）。

完整训练循环

因为每一步都是随机采样，本示例没有传统意义上的 epoch 概念。batch size 为 128，Keras 原示例大约训练 7000 步：

for step in range(7000): loss = jit_step() if step%100 == 0: acc = (model(X_test).argmax(axis=1) == Y_test).mean().item() print(f"step {step:4d}, loss {loss.item():.2f}, acc {acc*100.:.2f}%")

很快就能达到 98%，通常能到 99%：

step 0, loss 4.03, acc 71.43% step 100, loss 0.34, acc 93.86% step 200, loss 0.23, acc 95.97% step 300, loss 0.18, acc 96.32% step 400, loss 0.18, acc 96.76% step 500, loss 0.13, acc 97.46% step 600, loss 0.14, acc 97.45% step 700, loss 0.10, acc 97.27% step 800, loss 0.23, acc 97.49% step 900, loss 0.13, acc 97.51% step 1000, loss 0.13, acc 97.88% step 1100, loss 0.11, acc 97.72% step 1200, loss 0.14, acc 97.65% step 1300, loss 0.12, acc 98.04% step 1400, loss 0.25, acc 98.17% step 1500, loss 0.11, acc 97.86% step 1600, loss 0.21, acc 98.21% step 1700, loss 0.14, acc 98.34% ...

每 100 步打印一次 loss 和测试集准确率。注意测试准确率每次都会在 CPU/GPU 上对全部 10000 张测试图做一次前向，因此打印会引入一些额外时间。

进阶探索

tinygrad 是纯 Python 且代码量很小，你现在可以自由地把玩它。以下是几个值得动手的方向：

• 多 GPU（两行代码）：给示例加多 GPU 支持只需要两步——把模型.shard到所有 GPU 上（Tensor.shard，见 tinygrad/tensor.py，把参数按指定轴切分到多个设备），再把数据集按 batch 维度.shard。你能找到这两行加在哪里吗？
• 调试与代码查看：with Context(DEBUG=2)显示正在运行的 kernels，DEBUG=4显示生成的代码。所有Context变量都可以用环境变量替代（例如DEBUG=4、DEV=CUDA）。
• BEAM 搜索自动调优：with Context(BEAM=2)会对每个 kernel 做 beam 搜索，在多种候选实现中找出你硬件上跑得最快的版本。搜索完成后 tinygrad 通常能在速度上与 PyTorch 竞争，且结果会缓存（CACHELEVEL/IGNORE_BEAM_CACHE控制缓存行为，见 tinygrad/helpers.py），下次无需重新搜索。

参考：仓库中的完整 MNIST 示例

如果想看一个把上述所有技巧（@TinyJit装饰器、@Context(TRAINING=1)、@function、Tensor.sequential、GlobalCounters.reset()、trange进度条、getenv参数化）整合到更复杂 CNN 的成品，可以直接阅读 examples/beautiful_mnist.py。它支持用环境变量切换优化器（MUON=1、SGD=1、默认 Adam）、切换 Fashion-MNIST（FASHION=1）、调整 batch size（BS=512）与步数（STEPS=70），还内置了TARGET_EVAL_ACC_PCT验收断言，非常适合作为进一步实验的起点。相关测试可以参考 test/models/test_mnist.py 与 test/test_tiny.py。

从零训练一个 MNIST 模型是了解 tinygrad 设计哲学的最佳起点：普通 Python 类即模型、惰性求值、Context全局变量、以及把"整个训练步"整体 JIT 的TinyJit。掌握这些概念后，你就可以顺着 docs/quickstart.md 和 docs/tensor 系列文档继续深入，或者直接挑战 examples 里更复杂的模型了。

【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygrad