make_pair and first attempt at hlb_cifar10
ADD
examples/hlb_cifar10.py
+#!/usr/bin/env python3
+# tinygrad implementation of https://github.com/tysam-code/hlb-CIFAR10/blob/main/main.py
+# https://myrtle.ai/learn/how-to-train-your-resnet-8-bag-of-tricks/
+from datasets import fetch_cifar
+from tinygrad import nn
+from tinygrad.nn import optim
+from extra.training import train
+from extra.utils import get_parameters
+
+class ConvGroup:
+ def __init__(self, channels_in, channels_out, short):
+ self.short = short
+ self.conv = [nn.Conv2d(channels_in if i == 0 else channels_out, channels_out, kernel_size=3, padding=1, bias=False) for i in range(1 if short else 3)]
+ self.norm = [nn.BatchNorm2D(channels_out) for _ in range(1 if short else 3)]
+
+ if not self.short:
+ self.se1 = nn.Linear(channels_out, channels_out//16)
+ self.se2 = nn.Linear(channels_out//16, channels_out)
+
+ def __call__(self, x):
+ x = self.conv[0](x).max_pool2d(2)
+ x = self.norm[0](x).gelu()
+ if self.short: return x
+ residual = x
+ mult = self.se2(self.se1(residual.mean((2,3))).gelu()).sigmoid().reshape(x.shape[0], x.shape[1], 1, 1)
+ x = self.norm[1](self.conv[1](x)).gelu()
+ x = self.norm[2](self.conv[2](x) * mult).gelu()
+ return x + residual
+
+class SpeedyResNet:
+ def __init__(self):
+ # TODO: add whitening
+ self.net = [
+ nn.Conv2d(3, 64, kernel_size=1),
+ nn.BatchNorm2D(64),
+ lambda x: x.gelu(),
+ ConvGroup(64, 128, short=False),
+ ConvGroup(128, 256, short=True),
+ ConvGroup(256, 512, short=False),
+ lambda x: x.max((2,3)),
+ nn.Linear(512, 1000, bias=False)
+ ]
+
+ def __call__(self, x): return x.sequential(self.net)
+
+def train_cifar():
+ X,Y = fetch_cifar()
+ model = SpeedyResNet()
+ optimizer = optim.SGD(get_parameters(model))
+ train(model, X, Y, optimizer, steps=X.shape[0]//512, BS=512)
+
+if __name__ == "__main__":
+ train_cifar()

MODIFY
extra/training.py
extra/training.py
@@ -24,7 +24,7 @@ def train(model, X_train, Y_train, optim, steps, BS=128, lossfn=sparse_categoric
y = target_transform(Y_train[samp])
# network
- out = model.forward(x)
+ out = model.forward(x) if hasattr(model, 'forward') else model(x)
loss = lossfn(out, y)
optim.zero_grad()

MODIFY
test/test_ops.py
test/test_ops.py
@@ -421,7 +421,7 @@ class TestOps(unittest.TestCase):
lambda x: Tensor.max_pool2d(x, kernel_size=ksz))
def test_maxpool2d(self):
- for ksz in [(2,2), (3,3), (3,2), (5,5), (5,1)]:
+ for ksz in [(2,2), (3,3), 2, 3, (3,2), (5,5), (5,1)]:
with self.subTest(kernel_size=ksz):
helper_test_op([(32,2,110,28)],
lambda x: torch.nn.functional.max_pool2d(x, kernel_size=ksz),

MODIFY
tinygrad/helpers.py
tinygrad/helpers.py
@@ -9,6 +9,7 @@ def all_same(items): return all(x == items[0] for x in items) if len(items) > 0
def colored(st, color): return f"\u001b[{30+['black', 'red', 'green', 'yellow', 'blue', 'magenta', 'cyan', 'white'].index(color)}m{st}\u001b[0m" # replace the termcolor library with one line
def partition(lst, fxn): return [x for x in lst if fxn(x)], [x for x in lst if not fxn(x)]
def modn(x, a): return -((-x)%a) if x < 0 else x%a
+def make_pair(x): return (x,x) if isinstance(x, int) else x
def reduce_shape(shape, axis): return tuple(1 if i in axis else shape[i] for i in range(len(shape)))
def shape_to_axis(old_shape, new_shape):

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
from __future__ import annotations
import inspect, functools, importlib, itertools
import numpy as np
-from tinygrad.helpers import prod, argfix
+from tinygrad.helpers import prod, argfix, make_pair
from typing import List, Tuple, Callable, Optional
from tinygrad.lazy import Device, LazyBuffer
@@ -251,8 +251,8 @@ class Tensor:
xup = self[:, :, :self.shape[2]-self.shape[2]%py, :self.shape[3]-self.shape[3]%px] if (self.shape[2]%py != 0) or (self.shape[3]%px != 0) else self
return xup.reshape(shape=(xup.shape[0], xup.shape[1], xup.shape[2]//py, py, xup.shape[3]//px, px))
- def avg_pool2d(self, kernel_size=(2,2)): return self._pool2d(*kernel_size).mean(axis=(3,5))
- def max_pool2d(self, kernel_size=(2,2)): return self._pool2d(*kernel_size).max(axis=(3,5))
+ def avg_pool2d(self, kernel_size=(2,2)): return self._pool2d(*make_pair(kernel_size)).mean(axis=(3,5))
+ def max_pool2d(self, kernel_size=(2,2)): return self._pool2d(*make_pair(kernel_size)).max(axis=(3,5))
def conv2d(self, weight, bias=None, **kwargs):
ret = self._conv2d(weight, **kwargs)

make_pair and first attempt at hlb_cifar10