add reshape support and OMG the CONVS are SO SLOW
MODIFY
test/mnist.py
test/mnist.py
#!/usr/bin/env python
+import os
import numpy as np
from tinygrad.tensor import Tensor
from tinygrad.utils import layer_init_uniform, fetch_mnist
@@ -18,16 +19,35 @@ class TinyBobNet:
def forward(self, x):
return x.dot(self.l1).relu().dot(self.l2).logsoftmax()
+# create a model with a conv layer
+# perfect if you like slow speeds and very little accuracy gains
+class TinyConvNet:
+ def __init__(self):
+ self.chans = 4
+ self.c1 = Tensor(layer_init_uniform(self.chans,1,3,3))
+ self.l1 = Tensor(layer_init_uniform(26*26*self.chans, 128))
+ self.l2 = Tensor(layer_init_uniform(128, 10))
+
+ def forward(self, x):
+ x.data = x.data.reshape((-1, 1, 28, 28)) # hacks
+ x = x.conv2d(self.c1).reshape(Tensor(np.array((-1, 26*26*self.chans)))).relu()
+ return x.dot(self.l1).relu().dot(self.l2).logsoftmax()
+
+if os.getenv("CONV") == "1":
+ model = TinyConvNet()
+ optim = optim.Adam([model.c1, model.l1, model.l2], lr=0.001)
+ steps = 400
+else:
model = TinyBobNet()
optim = optim.SGD([model.l1, model.l2], lr=0.001)
-#optim = optim.Adam([model.l1, model.l2], lr=0.001)
+ steps = 1000
BS = 128
losses, accuracies = [], []
-for i in (t := trange(1000)):
+for i in (t := trange(steps)):
samp = np.random.randint(0, X_train.shape[0], size=(BS))
- x = Tensor(X_train[samp].reshape((-1, 28*28)))
+ x = Tensor(X_train[samp].reshape((-1, 28*28)).astype(np.float32))
Y = Y_train[samp]
y = np.zeros((len(samp),10), np.float32)
# correct loss for NLL, torch NLL loss returns one per row

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
@@ -36,6 +36,8 @@ class Tensor:
if len(self._ctx.parents) == 1:
grads = [grads]
for t,g in zip(self._ctx.parents, grads):
+ if g is None:
+ continue
if g.shape != t.data.shape:
print("grad shape must match tensor shape in %r, %r != %r" %
(self._ctx, g.shape, t.data.shape))
@@ -75,6 +77,18 @@ def register(name, fxn):
# **** implement a few functions ****
+class Reshape(Function):
+ @staticmethod
+ def forward(ctx, x, shape):
+ ctx.save_for_backward(x.shape)
+ return x.reshape(shape)
+
+ @staticmethod
+ def backward(ctx, grad_output):
+ in_shape, = ctx.saved_tensors
+ return grad_output.reshape(in_shape), None
+register('reshape', Reshape)
+
class Mul(Function):
@staticmethod
def forward(ctx, x, y):

add reshape support and OMG the CONVS are SO SLOW