Package: tinygrad-experiments
For something in between apytorchand akarpathy/microgradThis may not be the best deep learning framework, but it is a deep learning framework.The sub 1000 line core of it is intinygrad/Due to its extreme simplicity, it aims to be the easiest framework to add new accelerators to, with support for both inference and training. Support the simple basic ops, and you get SOTAvisionmodels/efficientnet.pyandlanguagemodels/transformer.pymodels.We are working on support for the Apple Neural Engine and the Google TPU in theaccel/folder. Eventually,we will build custom hardwarefor tinygrad, and it will be blindingly fast. Now, it is slow.Installationgitclonehttps://github.com/geohot/tinygrad.gitcdtinygrad
python3setup.pydevelopContributingThere's a lot of interest in tinygrad lately. Here's some guidelines for contributing:Bugfixes are the best and always welcome! Likethis one.If you don't understand the code you are changing, don't change it!All code golf PRs will be closed, butconceptual cleanupsare great.Features are welcome. Though if you are adding a feature, you need to include tests.Improving test coverage is great, with reliable non brittle tests.Examplefromtinygrad.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)# dz/dxprint(y.grad)# dz/dySame example in torchimporttorchx=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/dyNeural 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.nn.optim, write some boilerplate minibatching code, and you have all you need.Neural network example (from test/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).logsoftmax()model=TinyBobNet()optim=optim.SGD([model.l1,model.l2],lr=0.001)# ... and complete like pytorch, with (x,y) dataout=model.forward(x)loss=out.mul(y).mean()optim.zero_grad()loss.backward()optim.step()GPU and Accelerator Supporttinygrad supports GPUs through PyOpenCL.fromtinygrad.tensorimportTensor(Tensor.ones(4,4).gpu()+Tensor.ones(4,4).gpu()).cpu()ANE Support?! (broken)If all you want to do is ReLU, you are in luck! You can do very fast ReLU (at least 30 MEGAReLUs/sec confirmed)Requires your Python to be signed withane/lib/sign_python.shto add thecom.apple.ane.iokit-user-accessentitlement, which also requiressudo nvram boot-args="amfi_get_out_of_my_way=1 ipc_control_port_options=0". Build the library withane/lib/build.shIn order to set boot-args and for the AMFI kext to respect that arg, runcsrutil enable --without-kext --without-nvramin recovery mode.fromtinygrad.tensorimportTensora=Tensor([-2,-1,0,1,2]).ane()b=a.relu()print(b.cpu())Warning: do not rely on the ANE port. It segfaults sometimes. So if you were doing something important with tinygrad and wanted to use the ANE, you might have a bad time.hlops (in tensor.py)hlops are syntactic sugar around mlops. They support most things torch does.mlopsmlops are mid level ops, there's 15 of them. They understand memory allocation and derivativesRelu, Log, Exp                          # unary ops
Sum, Max                                # reduce ops (with axis argument)
Add, Sub, Mul, Pow                      # binary ops (no broadcasting, use expand)
Reshape, Permute, Slice, Expand, Flip   # movement ops
Conv2D(NCHW)                            # processing op (Matmul is also Conv2D)You no longer need to write mlops for a new acceleratorAdding an accelerator (llops)The autodiff stuff is all in mlops now so you can focus on the raw operationsBuffer                                                     # class of memory on this device
unary_op  (RELU, EXP, LOG, NEG, SIGN)                      # A -> A
reduce_op (SUM, MAX)                                       # A -> B (smaller size, B has 1 in shape)
binary_op (ADD, SUB, MUL, DIV, POW, CMPEQ)                 # A + B -> C (all the same size)
movement_op (RESHAPE, PERMUTE, PAD, SHRINK, EXPAND, FLIP)  # A -> B (different size)
processing_op (CONV)                                       # A + B -> CWhen tinygrad moves to lazy evaluation, optimizations will happen here.ImageNet inferenceDespite being tiny, tinygrad supports the full EfficientNet. Pass in a picture to discover what it is.ipython3examples/efficientnet.pyhttps://media.istockphoto.com/photos/hen-picture-id831791190Or, if you have a webcam and cv2 installedipython3examples/efficientnet.pywebcamPROTIP: Set "GPU=1" environment variable if you want this to go faster.PROPROTIP: Set "DEBUG=1" environment variable if you want to see why it's slow.tinygrad supports Stable Diffusion!RunTORCH=1 python3 examples/stable_diffusion.py(or without torch:OPT=2 OPENCL=1 python3 examples/stable_diffusion.py)"a horse sized cat eating a bagel"tinygrad supports GANsSeeexamples/mnist_gan.pytinygrad supports yoloSeeexamples/yolov3.pyThe promise of smalltinygrad will always be below 1000 lines. If it isn't, we will revert commits until tinygrad becomes smaller.Drawing Execution GraphNodes are TensorsBlack edge is a forward passBlue edge is a backward passRed edge is data the backward pass depends onPurple edge is intermediates created in the forwardGRAPH=1python3test/test_mnist.pyTestMNIST.test_sgd_onestep# requires dot, outputs /tmp/net.svgRunning testspython3-mpytest