tinygrad輕量級深度學習框架
tinygrad 是一個輕量級深度學習框架，它的目標是成為最容易添加新加速器的框架，同時支持推理和訓練。
tinygrad 支持簡單的基本操作，包含 SOTA 視覺 models/efficientnet.py和 語言模型 models/transformer.py。
tinygrad 正在努力支持 accel/ 文件夾中的 Apple Neural Engine 和 Google TPU。最終將為 tinygrad 構(gòu)建定制硬件，以達到終極目標：極速推理/訓練框架。（但現(xiàn)在它還很慢）
安裝
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
神經(jīng)網(wǎng)絡(luò)示例（來自 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()
評論
圖片
表情