这可能不是最好的深度学习框架,但它是一个深度学习框架。
由于其极其简单,它旨在成为添加新加速器的最简单框架,并支持推理和训练。如果XLA是CISC,tinygrad就是RISC。
Tinygrad仍然是Alpha软件,但我们筹集了一些资金来使其变得更好。总有一天,我们会把芯片流下来。
尝试使用矩阵。看看,尽管风格如此,但它如何融合成一个内核,具有懒惰的力量。
DEBUG=3 python3 -c "from tinygrad.tensor import Tensor;
N = 1024; a, b = Tensor.rand(N, N), Tensor.rand(N, N);
c = (a.reshape(N, 1, N) * b.permute(1,0).reshape(1, N, N)).sum(axis=2);
print((c.numpy() - (a.numpy() @ b.numpy())).mean())"
我们可以更改为查看生成的代码。
DEBUG
4
事实证明,神经网络所需的 90% 是一个不错的自动/张量库。加入一个优化器、一个数据加载器和一些计算,你就拥有了你需要的一切。
from tinygrad.tensor import Tensor
import tinygrad.nn.optim as optim
class TinyBobNet:
def __init__(self):
self.l1 = Tensor.uniform(784, 128)
self.l2 = Tensor.uniform(128, 10)
def forward(self, x):
return x.dot(self.l1).relu().dot(self.l2).log_softmax()
model = TinyBobNet()
optim = optim.SGD([model.l1, model.l2], lr=0.001)
# ... complete data loader here
out = model.forward(x)
loss = out.mul(y).mean()
optim.zero_grad()
loss.backward()
optim.step()
tinygrad已经支持许多加速器,包括:
而且很容易添加更多!你选择的加速器只需要支持总共 26 个(可选 27 个)低级操作。有关详细信息,请参阅有关添加新加速器的文档。
目前推荐的安装 tinygrad 的方法是从源代码开始。
git clone https://github.com/geohot/tinygrad.git
cd tinygrad
python3 -m pip install -e .
别忘了最后!
.
文档和快速入门指南可以在 docs/ 目录中找到。
from tinygrad.tensor 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很感兴趣。以下是一些贡献的基本准则:
其他指南可在 CONTRIBUTING.md 中找到。
有关如何运行完整测试套件的更多示例,请参阅 CI 工作流。
一些例子:
python3 -m pip install -e '.[testing]'
python3 -m pytest
python3 -m pytest -v -k TestTrain
python3 ./test/models/test_train.py TestTrain.test_efficientnet