tinygrad MNIST 入门教程从模型定义、训练到 TinyJit 提速的完整实战【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygradMNIST 手写数字识别是深度学习框架的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 githttps://github.com/tinygrad/tinygrad.gitColab 上不需要任何特殊的镜像但为了速度建议选择 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如paddingsame但要求stride1否则会抛出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 字节头。它还支持fashionTrue参数切换到 Fashion-MNIST 数据集以及deviceNone参数指定加载到哪个设备。同一文件里的cifar()则展示了如何用tar_extract解包 CIFAR-10 的 tar 压缩包可作为参考。模型推理惰性求值由于 MNIST 足够小mnist()会把数据集直接复制到默认设备上。于是创建模型 跑一次推理 算准确率只需要model Model() acc (model(X_test).argmax(axis1) Y_test).mean() # NOTE: tinygrad is lazy, and hasnt 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(TRAINING1) def step(): samples Tensor.randint(batch_size, highX_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(TRAINING1)把TRAINING这个全局上下文变量设为 1。从 tinygrad/helpers.py 可见TRAINING ContextVar(TRAINING, 0)而 tinygrad/nn/optim.py 的Optimizer.schedule_step会检查if not TRAINING: raise RuntimeError(...)——不开启 TRAINING 直接调用优化器会直接报错提示你使用Context(TRAINING1)。这也让nn.BatchNorm等层能区分训练/推理统计量。Tensor.randint(batch_size, highX_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参数默认reductionmean。.backward()自动微分构建并执行反向传播。optim.zero_grad()/optim.step()清空梯度、更新参数。计时单步import timeit timeit.repeat(step, repeat5, number1) #[0.08268719699981375, # 0.07478952900009972, # 0.07714716600003158, # 0.07785399599970333, # 0.07605237000007037]T4 Colab 上约 75 ms 一步。如果想看每个 kernel 的时间分解可以用from tinygrad import GlobalCounters, Context GlobalCounters.reset() with Context(DEBUG2): step()DEBUG是 tinygrad 最常用的调试开关之一tinygrad/helpers.py 中DEBUG ContextVar(DEBUG, 0)DEBUG2显示每个被运行的 kernel 及其耗时DEBUG4会进一步打印生成的 kernel 源码所有的Context变量同时也都能通过同名环境变量设置例如DEBUG2 python train.py。关于优化器仓库 tinygrad/nn/optim.py 按复杂度排序提供了完整家族——SGD本质是关闭 trust ratio 的LARS、Adam/AdamW本质是关闭 trust ratio 的LAMB、MuonLARS Newton-Schulz 迭代ns_steps5等。Adam默认lr0.001, b10.9, b20.999, eps1e-8优化器还会把学习率提升到至少 float32 精度存储。为什么这么慢——引入 TinyJit75 ms 一步对调试来说足够快但对训练来说并不理想。与 PyTorch 不同tinygrad 天生不是为这种慢速逐步执行设计的。这里引入 tinygrad 第一个标志性概念TinyJit。from tinygrad import TinyJit jit_step TinyJit(step)注意它也可以用作装饰器TinyJit。再计时import timeit timeit.repeat(jit_step, repeat5, number1) # [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__实现完整印证前两次调用正常执行并捕获第一次调用cnt0按原样运行第二次cnt1进入capture阶段用capturing上下文把执行过程中触发的所有 kernelOps.LINEAR序列记录下来然后经jit_lower做参数化输入 buffer 映射为 PARAM、内存规划memory_plan_rewrite和可选的 graph 化graph_split_rewrite受JIT_BATCH_SIZE控制默认 32最终打包成一个CapturedJit。从第三次起只做回放跳过整个 Python 执行过程只按捕获的 kernel 序列在 GPU/CPU 上重放Python 调度开销被完全消除。这就是dumb replay on different data——把 JIT 当作对不同数据的无脑重放。注意事项任何影响 kernel 的非 tinygrad Python 值都会在第二次运行后冻结因为回放不再经过 Python。不过Tensor的随机函数如Tensor.randint是编译进 kernel 的因此能正常工作——这也是上面step()里随机采样放在函数内部仍然有效的原因。与大多数框架不同tinygrad 的 JIT连优化器也一起 JIToptim.step()也在被捕获的函数内。需要理解Tensor输入的规范化约定_prepare_jit_inputs要求 JIT 函数的 Tensor 参数必须是真实 buffer虚拟视图会报JIT inputs must be real buffers; use .clone()且不允许重复输入。若捕获期间生成了惰性.grad等held_bufs收集逻辑会保证它们不被内存规划器误释放。另外_TinyJit还支持reset()重新捕获与pruneTrue剪枝掉与输出无关的一次性 kernel。完整训练循环因为每一步都是随机采样本示例没有传统意义上的 epoch 概念。batch size 为 128Keras 原示例大约训练 7000 步for step in range(7000): loss jit_step() if step%100 0: acc (model(X_test).argmax(axis1) Y_test).mean().item() print(fstep {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(DEBUG2)显示正在运行的 kernelsDEBUG4显示生成的代码。所有Context变量都可以用环境变量替代例如DEBUG4、DEVCUDA。BEAM 搜索自动调优with Context(BEAM2)会对每个 kernel 做 beam 搜索在多种候选实现中找出你硬件上跑得最快的版本。搜索完成后 tinygrad 通常能在速度上与 PyTorch 竞争且结果会缓存CACHELEVEL/IGNORE_BEAM_CACHE控制缓存行为见 tinygrad/helpers.py下次无需重新搜索。参考仓库中的完整 MNIST 示例如果想看一个把上述所有技巧TinyJit装饰器、Context(TRAINING1)、function、Tensor.sequential、GlobalCounters.reset()、trange进度条、getenv参数化整合到更复杂 CNN 的成品可以直接阅读 examples/beautiful_mnist.py。它支持用环境变量切换优化器MUON1、SGD1、默认 Adam、切换 Fashion-MNISTFASHION1、调整 batch sizeBS512与步数STEPS70还内置了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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 tinygrad MNIST 入门教程：从模型定义、训练到 TinyJit 提速的完整实战
tinygrad MNIST 入门教程从模型定义、训练到 TinyJit 提速的完整实战【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygradMNIST 手写数字识别是深度学习框架的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 githttps://github.com/tinygrad/tinygrad.gitColab 上不需要任何特殊的镜像但为了速度建议选择 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如paddingsame但要求stride1否则会抛出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 字节头。它还支持fashionTrue参数切换到 Fashion-MNIST 数据集以及deviceNone参数指定加载到哪个设备。同一文件里的cifar()则展示了如何用tar_extract解包 CIFAR-10 的 tar 压缩包可作为参考。模型推理惰性求值由于 MNIST 足够小mnist()会把数据集直接复制到默认设备上。于是创建模型 跑一次推理 算准确率只需要model Model() acc (model(X_test).argmax(axis1) Y_test).mean() # NOTE: tinygrad is lazy, and hasnt 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(TRAINING1) def step(): samples Tensor.randint(batch_size, highX_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(TRAINING1)把TRAINING这个全局上下文变量设为 1。从 tinygrad/helpers.py 可见TRAINING ContextVar(TRAINING, 0)而 tinygrad/nn/optim.py 的Optimizer.schedule_step会检查if not TRAINING: raise RuntimeError(...)——不开启 TRAINING 直接调用优化器会直接报错提示你使用Context(TRAINING1)。这也让nn.BatchNorm等层能区分训练/推理统计量。Tensor.randint(batch_size, highX_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参数默认reductionmean。.backward()自动微分构建并执行反向传播。optim.zero_grad()/optim.step()清空梯度、更新参数。计时单步import timeit timeit.repeat(step, repeat5, number1) #[0.08268719699981375, # 0.07478952900009972, # 0.07714716600003158, # 0.07785399599970333, # 0.07605237000007037]T4 Colab 上约 75 ms 一步。如果想看每个 kernel 的时间分解可以用from tinygrad import GlobalCounters, Context GlobalCounters.reset() with Context(DEBUG2): step()DEBUG是 tinygrad 最常用的调试开关之一tinygrad/helpers.py 中DEBUG ContextVar(DEBUG, 0)DEBUG2显示每个被运行的 kernel 及其耗时DEBUG4会进一步打印生成的 kernel 源码所有的Context变量同时也都能通过同名环境变量设置例如DEBUG2 python train.py。关于优化器仓库 tinygrad/nn/optim.py 按复杂度排序提供了完整家族——SGD本质是关闭 trust ratio 的LARS、Adam/AdamW本质是关闭 trust ratio 的LAMB、MuonLARS Newton-Schulz 迭代ns_steps5等。Adam默认lr0.001, b10.9, b20.999, eps1e-8优化器还会把学习率提升到至少 float32 精度存储。为什么这么慢——引入 TinyJit75 ms 一步对调试来说足够快但对训练来说并不理想。与 PyTorch 不同tinygrad 天生不是为这种慢速逐步执行设计的。这里引入 tinygrad 第一个标志性概念TinyJit。from tinygrad import TinyJit jit_step TinyJit(step)注意它也可以用作装饰器TinyJit。再计时import timeit timeit.repeat(jit_step, repeat5, number1) # [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__实现完整印证前两次调用正常执行并捕获第一次调用cnt0按原样运行第二次cnt1进入capture阶段用capturing上下文把执行过程中触发的所有 kernelOps.LINEAR序列记录下来然后经jit_lower做参数化输入 buffer 映射为 PARAM、内存规划memory_plan_rewrite和可选的 graph 化graph_split_rewrite受JIT_BATCH_SIZE控制默认 32最终打包成一个CapturedJit。从第三次起只做回放跳过整个 Python 执行过程只按捕获的 kernel 序列在 GPU/CPU 上重放Python 调度开销被完全消除。这就是dumb replay on different data——把 JIT 当作对不同数据的无脑重放。注意事项任何影响 kernel 的非 tinygrad Python 值都会在第二次运行后冻结因为回放不再经过 Python。不过Tensor的随机函数如Tensor.randint是编译进 kernel 的因此能正常工作——这也是上面step()里随机采样放在函数内部仍然有效的原因。与大多数框架不同tinygrad 的 JIT连优化器也一起 JIToptim.step()也在被捕获的函数内。需要理解Tensor输入的规范化约定_prepare_jit_inputs要求 JIT 函数的 Tensor 参数必须是真实 buffer虚拟视图会报JIT inputs must be real buffers; use .clone()且不允许重复输入。若捕获期间生成了惰性.grad等held_bufs收集逻辑会保证它们不被内存规划器误释放。另外_TinyJit还支持reset()重新捕获与pruneTrue剪枝掉与输出无关的一次性 kernel。完整训练循环因为每一步都是随机采样本示例没有传统意义上的 epoch 概念。batch size 为 128Keras 原示例大约训练 7000 步for step in range(7000): loss jit_step() if step%100 0: acc (model(X_test).argmax(axis1) Y_test).mean().item() print(fstep {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(DEBUG2)显示正在运行的 kernelsDEBUG4显示生成的代码。所有Context变量都可以用环境变量替代例如DEBUG4、DEVCUDA。BEAM 搜索自动调优with Context(BEAM2)会对每个 kernel 做 beam 搜索在多种候选实现中找出你硬件上跑得最快的版本。搜索完成后 tinygrad 通常能在速度上与 PyTorch 竞争且结果会缓存CACHELEVEL/IGNORE_BEAM_CACHE控制缓存行为见 tinygrad/helpers.py下次无需重新搜索。参考仓库中的完整 MNIST 示例如果想看一个把上述所有技巧TinyJit装饰器、Context(TRAINING1)、function、Tensor.sequential、GlobalCounters.reset()、trange进度条、getenv参数化整合到更复杂 CNN 的成品可以直接阅读 examples/beautiful_mnist.py。它支持用环境变量切换优化器MUON1、SGD1、默认 Adam、切换 Fashion-MNISTFASHION1、调整 batch sizeBS512与步数STEPS70还内置了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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考