tinygrad:不到1000行代码的深度学习框架,天才黑客开源GitHub 2.3k+ stars

清泛编译

近期,一个不到1000行的深度学习框架tinygrad火了,麻雀虽小,但五脏俱全,这个深度学习框架使用起来和PyTorch类似,目前已经开源在GitHub上,而且收获了2.3K星:

项目地址https://github.com/geohot/tinygrad

本质上说,tinygrad就是一个可微的Tensor库,但是对于训练一个简单的神经网络它是足够了,项目也给一个简单的case:

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()

另外tinygrad也支持GPU:

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

其实我们当然不会用这样的框架来真正训练网络,但是1000行的代码量我们倒是可以稍微啃一下,对理解深度学习框架内部原理还是非常有帮助的,希望感兴趣的小伙伴们可以查阅。

另外,一个类似的项目是micrograd (https://github.com/karpathy/micrograd),也可以比较一下。

tinygrad 深度学习

分享到:
评论加载中,请稍后...
创APP如搭积木 - 创意无限,梦想即时!
回到顶部