tinygrad 是一个轻量级深度学习框架，它的目标是成为最容易添加新加速器的框架，同时支持推理和训练。
tinygrad 支持简单的基本操作，包含SOTA 视觉 models/efficientnet.py和 语言模型 models/transformer.py。
tinygrad 正在努力支持 accel/ 文件夹中的AppleNeuralEngine和GoogleTPU。最终将为tinygrad构建定制硬件，以达到终极目标：极速推理/训练框架。（但现在它还很慢）
安装gitclonehttps://github.com/geohot/tinygrad.gitcdtinygradpython3-mpipinstall-e.例子fromtinygrad.tensorimportTensorx=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/dxprint(y.grad.numpy())#dz/dytorch中的相同示例importtorchx=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神经网络示例（来自test/models/test_mnist.py）fromtinygrad.tensorimportTensorimporttinygrad.nn.optimasoptimclassTinyBobNet:def__init__(self):self.l1=Tensor.uniform(784,128)self.l2=Tensor.uniform(128,10)defforward(self,x):returnx.dot(self.l1).relu().dot(self.l2).log_softmax()model=TinyBobNet()optim=optim.SGD([model.l1,model.l2],lr=0.001)#...andcompletelikepytorch,with(x,y)dataout=model.forward(x)loss=out.mul(y).mean()optim.zero_grad()loss.backward()optim.step()GPU和加速器支持tinygrad通过PyOpenCL支持GPU。fromtinygrad.tensorimportTensor(Tensor.ones(4,4).gpu()+Tensor.ones(4,4).gpu()).cpu()
评论