tinygrad：简约而强大的深度学习框架-CSDN博客

本文链接：https://blog.csdn.net/weixin_36829761/article/details/140510431

在当今的深度学习领域，框架如 PyTorch 和 karpathy/micrograd 已经广泛应用。然而，对于那些需要一个更简单、更易扩展的框架的开发者来说，tinygrad 提供了一个新的选择。tinygrad 由 tiny corp 维护，它致力于成为一个既支持推理又支持训练的深度学习框架。

设计理念

极简主义与易扩展性

tinygrad 的设计理念是极简主义。与 XLA 类比，如果 XLA 是复杂指令集计算 (CISC)，那么 tinygrad 就是精简指令集计算 (RISC)。这种简约的设计使得它成为添加新加速器最容易的框架之一。通过简化框架的架构，开发者可以更轻松地理解和扩展它。

早期发展与资金支持

尽管 tinygrad 仍处于 alpha 阶段，但它已经获得了一定的资金支持，以便进一步优化和发展。例如，tiny corp 已经筹集了 500 万美元，未来甚至计划设计专用芯片。

功能特性

支持 LLaMA 和稳定扩散

tinygrad 已经能够运行 LLaMA 和 Stable Diffusion，这展示了它在处理复杂模型方面的能力。

惰性计算

tinygrad 的一大特色是其惰性计算机制。以下示例展示了如何利用惰性计算进行矩阵乘法，并生成优化的计算内核：

DEBUG=3 python3 -c "from tinygrad import Tensor;
N = 1024; a, b = Tensor.rand(N, N), Tensor.rand(N, N);
c = (a.reshape(N, 1, N) * b.T.reshape(1, N, N)).sum(axis=2);
print((c.numpy() - (a.numpy() @ b.numpy())).mean())"

将 DEBUG 设置为 4 可以查看生成的代码。

构建神经网络

tinygrad 提供了构建神经网络所需的几乎所有基础功能，包括自动求导、张量库、优化器和数据加载器。以下是一个简单的神经网络示例：

from tinygrad import Tensor, nn

class LinearNet:
  def __init__(self):
    self.l1 = Tensor.kaiming_uniform(784, 128)
    self.l2 = Tensor.kaiming_uniform(128, 10)
  def __call__(self, x:Tensor) -> Tensor:
    return x.flatten(1).dot(self.l1).relu().dot(self.l2)

model = LinearNet()
optim = nn.optim.Adam([model.l1, model.l2], lr=0.001)

x, y = Tensor.rand(4, 1, 28, 28), Tensor([2,4,3,7])  # replace with real mnist dataloader

with Tensor.train():
  for i in range(10):
    optim.zero_grad()
    loss = model(x).sparse_categorical_crossentropy(y).backward()
    optim.step()
    print(i, loss.item())

完整的版本可以在 examples/beautiful_mnist.py 中找到，该版本在约 5 秒内获得 98% 的准确率。

支持的加速器

tinygrad 目前已经支持多种加速器，包括：

添加新的加速器也非常容易，只需要支持大约 25 个低级操作即可。

安装

目前推荐从源代码安装 tinygrad。

从源代码安装

git clone https://github.com/tinygrad/tinygrad.git
cd tinygrad
python3 -m pip install -e .

直接安装（master 分支）

python3 -m pip install git+https://github.com/tinygrad/tinygrad.git

文档

详细的文档和快速入门指南可以在 docs 网站上找到，这些文档来自 docs/ 目录。

与 PyTorch 的简单比较

以下示例展示了 tinygrad 和 PyTorch 的相似性：

from tinygrad import Tensor

x = Tensor.eye(3, requires_grad=True)
y = Tensor([[2.0,0,-2.0]], requires_grad=True)
z = y.matmul(x).sum()
z.backward()

print(x.grad.numpy())  # dz/dx
print(y.grad.numpy())  # dz/dy

在 PyTorch 中实现相同的功能：

import torch

x = torch.eye(3, requires_grad=True)
y = torch.tensor([[2.0,0,-2.0]], requires_grad=True)
z = y.matmul(x).sum()
z.backward()

print(x.grad.numpy())  # dz/dx
print(y.grad.numpy())  # dz/dy

贡献指南

tinygrad 最近受到了很多关注。以下是一些有助于您的 PR 被接受的指南：

避免的事项

不允许“代码高尔夫”。虽然低行数是本项目的指导原则之一，但任何看起来像代码高尔夫的东西都会被关闭。真正的目标是减少复杂性和提高可读性，删除 \n 并不能帮助实现这一点。
除非您是知名贡献者，否则所有文档和空白字符的更改都会被关闭。编写文档的人应该是最了解代码库的人，那些没有证明自己了解代码库的人不应该修改文档。空白字符的更改既无用又有可能引入错误。
任何声称是“加速”的更改必须经过基准测试。总体目标是简化，即使您的 PR 使某些东西稍微快了一点，也必须权衡可维护性和可读性。
一般来说，tinygrad/ 文件夹外的代码没有经过严格测试，除非当前代码有问题，否则不应该更改它。
如果您的 PR 看起来“复杂”、差异很大或增加了很多行，它将不会被审查或合并。考虑将其拆分为较小的 PR，每个 PR 都是单独的明显改进。常见的模式是，在添加新功能之前进行必要的重构。如果您可以（干净地）重构到该功能仅需 3 行更改，这将是一个很好的开始，并且易于我们审查。

我们希望的贡献

Bug 修复（带有回归测试）是非常受欢迎的！这个库还没有到 1.0，因此如果您遇到 bug，请修复它，编写测试并提交 PR，这是非常有价值的工作。
解决悬赏任务！tinygrad 提供现金奖励以鼓励对库进行某些改进。所有新代码都应具有高质量并经过充分测试。
功能添加。然而，如果您要添加功能，请考虑行数权衡。如果是 3 行代码，它需要满足的有用性门槛比 30 或 300 行代码更低。所有功能必须有回归测试。在没有其他约束的情况下，您的功能 API 应与 torch 或 numpy 保持一致。
明显的重构。在一般情况下，如果您的重构没有明显的优势，它将被关闭。但有些重构是非常棒的！要考虑深层次的可读性。一个空白字符的更改或移动几个函数是无用的，但如果您意识到两个 100 行的函数实际上可以使用同一个 110 行的带参数函数，同时提高可读性，这是一个很大的胜利。重构应通过过程重播。
测试/模糊测试。如果您能添加非脆弱的测试，我们非常欢迎。我们这里也有一些模糊测试，通过改进它们可以发现大量的 bug。发现 bug，甚至编写预期失败的测试（@unittest.expectedFailure）也是很好的。这是我们进步的方式。
删除核心 tinygrad/ 文件夹中的死代码。我们不关心 extra 中的代码，但删除核心库中的死代码是很好的。这样新的开发者就有更少的代码需要阅读和困惑。

运行测试

您应该使用 pre-commit install 安装预提交钩子。这将在线性器、mypy 和子集测试在每次提交时运行。

有关运行完整测试套件的更多示例，请参阅 CI 工作流。

一些在本地运行测试的示例：

python3 -m pip install -e '.[testing]'  # 安装测试所需的额外依赖
python3 test/test_ops.py                # 仅运行 ops 测试
python3 -m pytest test/                 # 运行整个测试套件