tinygrad轻量级深度学习框架
tinygrad 是一个轻量级深度学习框架，它的目标是成为最容易添加新加速器的框架，同时支持推理和训练。
tinygrad 支持简单的基本操作，包含 SOTA 视觉 models/efficientnet.py和 语言模型 models/transformer.py。
tinygrad 正在努力支持 accel/ 文件夹中的 Apple Neural Engine 和 Google TPU。最终将为 tinygrad 构建定制硬件，以达到终极目标：极速推理/训练框架。（但现在它还很慢）
安装
git clone https://github.com/geohot/tinygrad.git cd tinygrad python3 -m pip install -e .
例子
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
torch 中的相同示例
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) # dz/dx print(y.grad) # dz/dy
神经网络示例（来自 test/models/test_mnist.py）
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) # ... and complete like pytorch, with (x,y) data out = model.forward(x) loss = out.mul(y).mean() optim.zero_grad() loss.backward() optim.step()
GPU 和加速器支持
tinygrad 通过 PyOpenCL 支持 GPU。
from tinygrad.tensor import Tensor (Tensor.ones(4,4).gpu() + Tensor.ones(4,4).gpu()).cpu()
评论