小型深度學(xué)習(xí)框架 | TinyGrad，不到1K行代碼（附代碼下載）

最近，天才黑客 George Hotz 開源了一個(gè)小型深度學(xué)習(xí)框架 tinygrad，兼具 PyTorch 和 micrograd 的功能。tinygrad 的代碼數(shù)量不到 1000 行，目前該項(xiàng)目獲得了 GitHub 1400 星。

在深度學(xué)習(xí)時(shí)代，谷歌、Facebook、百度等科技巨頭開源了多款框架來幫助開發(fā)者更輕松地學(xué)習(xí)、構(gòu)建和訓(xùn)練不同類型的神經(jīng)網(wǎng)絡(luò)。而這些大公司也花費(fèi)了很大的精力來維護(hù) TensorFlow、PyTorch 這樣龐大的深度學(xué)習(xí)框架。除了這類主流框架之外，開發(fā)者們也會(huì)開源一些小而精的框架或者庫(kù)。比如今年 4 月份，特斯拉人工智能部門主管 Andrej Karpathy 開源了其編寫的微型 autograd 引擎 micrograd，該引擎還用 50 行代碼實(shí)現(xiàn)了一個(gè)類 PyTorch api 的神經(jīng)網(wǎng)絡(luò)庫(kù)。目前，micrograd 項(xiàng)目的 GitHub star 量達(dá)到 1200 星。不久前，天才黑客 George Hotz（喬治 · 霍茲）開源了一個(gè)小型 Autograd Tensor 庫(kù) tinygrad，它介于 PyTorch 和 micrograd 之間，能夠滿足做深度學(xué)習(xí)的大部分要求。上線不到一個(gè)月，該項(xiàng)目在 GitHub 上已經(jīng)獲得 1400 星。

根據(jù) GitHub 內(nèi)容，下文對(duì) tinygrad 的安裝與使用做了簡(jiǎn)要介紹。感興趣的同學(xué)也可通過 George Hotz 的 YouTube 視頻進(jìn)行學(xué)習(xí)。

視頻地址：https://www.youtube.com/channel/UCwgKmJM4ZJQRJ-U5NjvR2dg

tinygrad 的安裝與使用

「tinygrad 可能不是最好的深度學(xué)習(xí)框架，但它確實(shí)是深度學(xué)習(xí)框架?！?/span>

George 在項(xiàng)目中保證，tinygrad 代碼量會(huì)永遠(yuǎn)小于 1000 行。

安裝
tinygrad 的安裝過程非常簡(jiǎn)單，只需使用以下命令：

pip3 install tinygrad --upgrade

from tinygrad.tensor import Tensor
x = Tensor.eye(3)y = Tensor([[2.0,0,-2.0]])z = y.matmul(x).sum()z.backward()
print(x.grad)  # dz/dxprint(y.grad)  # 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/dxprint(y.grad)  # dz/dy

滿足對(duì)神經(jīng)網(wǎng)絡(luò)的需求

一個(gè)不錯(cuò)的autograd張量庫(kù)可以滿足你對(duì)神經(jīng)網(wǎng)絡(luò) 90％的需求。從 tinygrad.optim 添加優(yōu)化器（SGD、RMSprop、Adam），再編寫一些 minibatching 樣板代碼，就可以實(shí)現(xiàn)你的需求。

示例如下：

from tinygrad.tensor import Tensorimport tinygrad.optim as optimfrom tinygrad.utils import layer_init_uniform
class TinyBobNet:  def __init__(self):    self.l1 = Tensor(layer_init_uniform(784, 128))    self.l2 = Tensor(layer_init_uniform(128, 10))
  def forward(self, x):    return x.dot(self.l1).relu().dot(self.l2).logsoftmax()
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()loss.backward()optim.step()

支持 GPUtinygrad 通過 PyOpenCL 支持 GPU。但后向傳播暫時(shí)無法支持所有 ops。

from tinygrad.tensor import Tensor(Tensor.ones(4,4).cuda() + Tensor.ones(4,4).cuda()).cpu()

ImageNet inference
「麻雀雖小，五臟俱全。」tinygrad 還能夠支持 full EfficientNet，輸入一張圖像，即可得到其類別。

ipython3 examples/efficientnet.py https://upload.wikimedia.org/wikipedia/commons/4/41/Chicken.jpg

如果你安裝了 webcam 和 cv2，則可以使用以下代碼：

ipython3 examples/efficientnet.py webcam

python -m pytest

• Train an EfficientNet on ImageNet

• Make broadcasting work on the backward pass (simple please)

• EfficientNet backward pass

• Tensors on GPU (a few more backward)

• Add a language model. BERT?

• Add a detection model. EfficientDet?

• Reduce code

• Increase speed

• Add features

／End.