October 24, 2020

305 words 2 mins read



You like pytorch? You like micrograd? You love tinygrad!

repo name geohot/tinygrad
repo link https://github.com/geohot/tinygrad
language Python
size (curr.) 176 kB
stars (curr.) 342
created 2020-10-18
license MIT License

Unit Tests

For something in between a pytorch and a karpathy/micrograd

This may not be the best deep learning framework, but it is a deep learning framework.

The Tensor class is a wrapper around a numpy array, except it does Tensor things.


pip3 install tinygrad


from tinygrad.tensor import Tensor

x = Tensor.eye(3)
y = Tensor([[2.0,0,-2.0]])
z = y.matmul(x).sum()

print(x.grad)  # dz/dx
print(y.grad)  # dz/dy

Same example in 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()

print(x.grad)  # dz/dx
print(y.grad)  # dz/dy

Neural networks?

It turns out, a decent autograd tensor library is 90% of what you need for neural networks. Add an optimizer (SGD, RMSprop, and Adam implemented) from tinygrad.optim, write some boilerplate minibatching code, and you have all you need.

Neural network example (from test/test_mnist.py)

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

GPU Support?!

tinygrad supports GPUs through PyOpenCL. Not all ops are supported yet.

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

ImageNet inference (on the micrograd puppy)

python3 examples/efficientnet.py

The promise of small

tinygrad will always be below 1000 lines. If it isn’t, we will revert commits until tinygrad becomes smaller.

Running tests

python -m pytest


  • Train an EfficientNet on ImageNet
    • Make broadcasting work on the backward pass (simple please)
    • EfficientNet backward pass
    • Tensors on GPU (GPU support, must support Mac)
  • Reduce code
  • Increase speed
  • Add features
comments powered by Disqus