在深度学习时代,谷歌、Facebook、百度等科技巨头开源了多款框架来帮助开发者更轻松地学习、构建和训练不同类型的神经网络。而这些大公司也花费了很大的精力来维护 TensorFlow、PyTorch这样庞大的深度学习框架。除了这类主流框架之外,开发者们也会开源一些小而精的框架或者库。比如今年4月份,特斯拉人工智能部门主管Andrej Karpathy开源了其编写的微型autograd引擎micrograd,该引擎还用50行代码实现了一个类PyTorch api的神经网络库。目前,micrograd项目的GitHub star量达到1200星。不久前,天才黑客George Hotz(乔治·霍兹)开源了一个小型Autograd Tensor库tinygrad,它介于PyTorch和micrograd之间,能够满足做深度学习的大部分要求。上线不到一个月,该项目在GitHub上已经获得1400星。
根据 GitHub 内容,下文对 tinygrad 的安装与使用做了简要介绍。感兴趣的同学也可通过 George Hotz 的 YouTube 视频进行学习。
视频地址:https://www.youtube.com/channel/UCwgKmJM4ZJQRJ-U5NjvR2dg
tinygrad 的安装与使用
「tinygrad 可能不是最好的深度学习框架,但它确实是深度学习框架。」
George 在项目中保证,tinygrad 代码量会永远小于 1000 行。
安装
tinygrad 的安装过程非常简单,只需使用以下命令:
pip3 install tinygrad --upgrade
示例
安装好 tinygrad 之后,就可以进行示例运行,代码如下:
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/dx
print(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/dx
print(y.grad) # dz/dy
满足对神经网络的需求
一个不错的autograd张量库可以满足你对神经网络 90%的需求。从 tinygrad.optim 添加优化器(SGD、RMSprop、Adam),再编写一些 minibatching 样板代码,就可以实现你的需求。
示例如下:
from tinygrad.tensor import Tensor
import tinygrad.optim as optim
from 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()
支持 GPU
tinygrad 通过 PyOpenCL 支持 GPU。但后向传播暂时无法支持所有 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
注意:如果你想加速运行,设置 GPU=1。
测试
运行以下代码可执行测试:
python -m pytest
此外,乔治 · 霍兹还计划添加语言模型、检测模型,进一步减少代码量、提升速度等。
TODO
- 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