trainer works with CIFAR
MODIFY
examples/train_efficientnet.py
examples/train_efficientnet.py
@@ -3,31 +3,64 @@ import time
import numpy as np
from extra.efficientnet import EfficientNet
from tinygrad.tensor import Tensor
-from tinygrad.utils import get_parameters
+from tinygrad.utils import get_parameters, fetch
+from tinygrad.utils import layer_init_uniform
+from tqdm import trange
import tinygrad.optim as optim
+import io
+import tarfile
+import pickle
+
+class TinyConvNet:
+ def __init__(self, classes=10):
+ conv = 3
+ inter_chan, out_chan = 8, 16 # for speed
+ self.c1 = Tensor(layer_init_uniform(inter_chan,3,conv,conv))
+ self.c2 = Tensor(layer_init_uniform(out_chan,inter_chan,conv,conv))
+ self.l1 = Tensor(layer_init_uniform(out_chan*6*6, classes))
+
+ def forward(self, x):
+ x = x.conv2d(self.c1).relu().max_pool2d()
+ x = x.conv2d(self.c2).relu().max_pool2d()
+ x = x.reshape(shape=[x.shape[0], -1])
+ return x.dot(self.l1).logsoftmax()
+
+def load_cifar():
+ tt = tarfile.open(fileobj=io.BytesIO(fetch('https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz')), mode='r:gz')
+ db = pickle.load(tt.extractfile('cifar-10-batches-py/data_batch_1'), encoding="bytes")
+ X = db[b'data'].reshape((-1, 3, 32, 32))
+ Y = np.array(db[b'labels'])
+ return X, Y
if __name__ == "__main__":
+ X_train, Y_train = load_cifar()
+ classes = 10
+
Tensor.default_gpu = os.getenv("GPU") is not None
- model = EfficientNet(int(os.getenv("NUM", "0")))
+ TINY = os.getenv("TINY") is not None
+ if TINY:
+ model = TinyConvNet(classes)
+ else:
+ model = EfficientNet(int(os.getenv("NUM", "0")), classes)
parameters = get_parameters(model)
- print(len(parameters))
+ print("parameters", len(parameters))
optimizer = optim.Adam(parameters, lr=0.001)
- BS = 16
- img = np.zeros((BS,3,224,224), dtype=np.float32)
+ #BS, steps = 16, 32
+ BS, steps = 64 if TINY else 16, 1024
- for i in range(32):
- print("running batch %d, %d tensors allocated" % (i, Tensor.allocated))
+ for i in (t := trange(steps)):
+ samp = np.random.randint(0, X_train.shape[0], size=(BS))
+
+ img = X_train[samp].astype(np.float32)
st = time.time()
out = model.forward(Tensor(img))
- et = time.time()
- print("forward %.2f s" % (et-st))
-
- Y = [0]*BS
+ fp_time = (time.time()-st)*1000.0
- y = np.zeros((BS,1000), np.float32)
- y[range(y.shape[0]),Y] = -1000.0
+ Y = Y_train[samp]
+ y = np.zeros((BS,classes), np.float32)
+ y[range(y.shape[0]),Y] = -classes
y = Tensor(y)
loss = out.logsoftmax().mul(y).mean()
@@ -35,13 +68,18 @@ if __name__ == "__main__":
st = time.time()
loss.backward()
- et = time.time()
- print("backward %.2f s" % (et-st))
+ bp_time = (time.time()-st)*1000.0
st = time.time()
optimizer.step()
- et = time.time()
- print("optimizer %.2f s" % (et-st))
+ opt_time = (time.time()-st)*1000.0
+
+ cat = np.argmax(out.cpu().data, axis=1)
+ accuracy = (cat == Y).mean()
+
+ # printing
+ t.set_description("loss %.2f accuracy %.2f -- %.2f %.2f %.2f -- %d" %
+ (loss.cpu().data, accuracy, fp_time, bp_time, opt_time, Tensor.allocated))
del out, y, loss

MODIFY
extra/efficientnet.py
extra/efficientnet.py
@@ -116,7 +116,7 @@ class MBConvBlock:
return x
class EfficientNet:
- def __init__(self, number=0):
+ def __init__(self, number=0, classes=1000):
self.number = number
global_params = [
# width, depth
@@ -171,8 +171,8 @@ class EfficientNet:
out_channels = round_filters(1280)
self._conv_head = Tensor.zeros(out_channels, in_channels, 1, 1)
self._bn1 = BatchNorm2D(out_channels)
- self._fc = Tensor.zeros(out_channels, 1000)
- self._fc_bias = Tensor.zeros(1000)
+ self._fc = Tensor.zeros(out_channels, classes)
+ self._fc_bias = Tensor.zeros(classes)
def forward(self, x):
x = x.pad2d(padding=(0,1,0,1))

trainer works with CIFAR