专栏中心

EEPW首页 > 专栏 > 小型深度学习框架 | TinyGrad,不到1K行代码(附代码下载)

小型深度学习框架 | TinyGrad,不到1K行代码(附代码下载)

发布人:CV研究院 时间:2022-12-22 来源:工程师 发布文章

最近,天才黑客 George Hotz 开源了一个小型深度学习框架 tinygrad,兼具 PyTorch  micrograd 的功能。tinygrad 的代码数量不到 1000 行,目前该项目获得了 GitHub 1400 星。

图片在深度学习时代,谷歌、Facebook、百度等科技巨头开源了多款框架来帮助开发者更轻松地学习、构建和训练不同类型的神经网络。而这些大公司也花费了很大的精力来维护 TensorFlowPyTorch 这样庞大的深度学习框架。除了这类主流框架之外,开发者们也会开源一些小而精的框架或者库。比如今年 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/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


满足对神经网络的需求


一个不错的autograd张量库可以满足你对神经网络 90%的需求。从 tinygrad.optim 添加优化器(SGD、RMSprop、Adam),再编写一些 minibatching 样板代码,就可以实现你的需求。

示例如下:






















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。但后向传播暂时无法支持所有 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

/End.


专栏文章内容及配图由作者撰写发布,仅供工程师学习之用,如有侵权或者其他违规问题,请联系本站处理。 联系我们

关键词: AI

相关推荐

AI聊天机器人能像医生一样推理吗?

存储器转型AI战略资源 台厂受惠

网络与存储 2026-05-19

基于Microchip MCU的AI/ML培训教程2

视频 2025-11-12

电子元件培训教材

思科凭借通用商用芯片与光模块赢得 AI 领域客户

AI/HPC新世代 COUPE光互连扮要角

网络与存储 2026-05-15

被动元件新周期:AI时代高端化、服务器化重构MLCC产业格局

前Qwen负责人林俊旸创业,目标融资规模为数亿美元

2026-05-14

基于Microchip MCU的AI/ML培训教程1

视频 2025-11-12

重新构想AI电源:塑造AI加速的未来(第三部分)

EEPW2018年6月刊(5G)

资源下载 2018-06-11

基于Microchip MCU的AI/ML培训教程3

视频 2025-11-12

PowiGaN for AI Data Centers: Unmatched Power Density and Reliability

视频 2025-12-19

联发科加速AI在地化应用布局

智能计算 2026-05-19

CSR8670CSR8675智能语音Alexa蓝牙方案开发

资源下载 2017-12-14

研华科技与Axelera AI深化战略合作 加速推动基于Europa平台的边缘AI创新

EEPW2018年3月刊(工业物联网)

尼吉康的事业介绍

视频 2025-07-25

国家“算力网”:像用水用电一样用AI

2026-05-18
更多 培训课堂
更多 焦点
更多 视频

技术专区