move good models to models, add a training step test
MODIFY
examples/efficientnet.py
examples/efficientnet.py
@@ -10,7 +10,7 @@ import numpy as np
np.set_printoptions(suppress=True)
from tinygrad.tensor import Tensor
from extra.utils import fetch, get_parameters
-from extra.efficientnet import EfficientNet
+from models.efficientnet import EfficientNet
def infer(model, img):
# preprocess image

MODIFY
examples/transformer.py
examples/transformer.py
@@ -6,7 +6,7 @@ import random
from tinygrad.tensor import Device
from extra.utils import get_parameters
from extra.training import train, evaluate
-from extra.transformer import Transformer
+from models.transformer import Transformer
from tinygrad.optim import Adam
# dataset idea from https://github.com/karpathy/minGPT/blob/master/play_math.ipynb

RENAME
extra/efficientnet.py
models/efficientnet.py
RENAME
extra/transformer.py
models/transformer.py
@@ -53,12 +53,12 @@ class TransformerBlock:
return x.reshape(shape=(bs, -1, embed_dim))
class Transformer:
- # L = cnt, H = embed_dim, A = num_heads
- def __init__(self, syms, maxlen, cnt, embed_dim, num_heads):
+ # L = layers, H = embed_dim, A = num_heads
+ def __init__(self, syms, maxlen, layers, embed_dim, num_heads):
self.maxlen, self.syms = maxlen, syms
self.embed = Tensor.uniform(maxlen+syms, embed_dim, requires_grad=False)
self.tbs = []
- for i in range(cnt):
+ for i in range(layers):
self.tbs.append(TransformerBlock(embed_dim, num_heads))
self.final = Tensor.uniform(embed_dim, syms)

ADD
test/test_train.py
+import unittest
+import time
+import tinygrad.optim as optim
+import numpy as np
+from tinygrad.tensor import Tensor
+from extra.training import train
+from extra.utils import get_parameters
+from models.efficientnet import EfficientNet
+from models.transformer import Transformer
+
+BS = 4
+
+def train_one_step(model,X,Y):
+ params = get_parameters(model)
+ pcount = 0
+ for p in params:
+ pcount += np.prod(p.shape)
+ optimizer = optim.Adam(params, lr=0.001)
+ print("stepping %r with %.1fM params bs %d" % (type(model), pcount/1e6, BS))
+ st = time.time()
+ train(model, X, Y, optimizer, steps=1, BS=BS)
+ et = time.time()-st
+ print("done in %.2f ms" % (et*1000.))
+
+class TestTrain(unittest.TestCase):
+ def test_efficientnet(self):
+ model = EfficientNet(0)
+ X = np.zeros((BS,3,224,224), dtype=np.float32)
+ Y = np.zeros((BS), dtype=np.int32)
+ train_one_step(model,X,Y)
+
+ def test_transformer(self):
+ # this should be small GPT-2, but the param count is wrong
+ model = Transformer(syms=10, maxlen=6, layers=12, embed_dim=768, num_heads=12)
+ X = np.zeros((BS,6), dtype=np.float32)
+ Y = np.zeros((BS,6), dtype=np.int32)
+ train_one_step(model,X,Y)
+
+ # these next two should be the mlperf models
+ def test_resnet(self):
+ # TODO: write this
+ pass
+
+ def test_bert(self):
+ # TODO: write this
+ pass
+
+if __name__ == '__main__':
+ unittest.main()

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
@@ -175,7 +175,7 @@ class Tensor:
if data.dtype != np.float32 and not Tensor.did_float_warning:
# warning? float64 is actually needed for numerical jacobian
- print(f"warning, {data.shape!r} isn't float32")
+ print(f"warning, {data.shape!r} isn't float32, it's {data.dtype}")
Tensor.did_float_warning = True
if device == Device.GPU:

move good models to models, add a training step test