add mnist example

mnist.py

@@ -0,0 +1,78 @@
+#!/usr/bin/env python
+import numpy as np
+from tensor import Tensor
+from tqdm import trange
+
+# load the mnist dataset
+
+def fetch(url):
+  import requests, gzip, os, hashlib, numpy
+  fp = os.path.join("/tmp", hashlib.md5(url.encode('utf-8')).hexdigest())
+  if not os.path.isfile(fp):
+    with open(fp, "rb") as f:
+      dat = f.read()
+  else:
+    with open(fp, "wb") as f:
+      dat = requests.get(url).content
+      f.write(dat)
+  return numpy.frombuffer(gzip.decompress(dat), dtype=np.uint8).copy()
+X_train = fetch("http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz")[0x10:].reshape((-1, 28, 28))
+Y_train = fetch("http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz")[8:]
+X_test = fetch("http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz")[0x10:].reshape((-1, 28, 28))
+Y_test = fetch("http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz")[8:]
+
+# train a model
+
+def layer_init(m, h):
+  ret = np.random.uniform(-1., 1., size=(m,h))/np.sqrt(m*h)
+  return ret.astype(np.float32)
+
+l1 = Tensor(layer_init(784, 128))
+l2 = Tensor(layer_init(128, 10))
+
+lr = 0.01
+BS = 128
+losses, accuracies = [], []
+for i in (t := trange(1000)):
+  samp = np.random.randint(0, X_train.shape[0], size=(BS))
+  
+  x = Tensor(X_train[samp].reshape((-1, 28*28)))
+  Y = Y_train[samp]
+  y = np.zeros((len(samp),10), np.float32)
+  y[range(y.shape[0]),Y] = -1.0
+  y = Tensor(y)
+  
+  x = x.dot(l1)
+  x = x.relu()
+  x = x_l2 = x.dot(l2)
+  x = x.logsoftmax()
+  x = x.mul(y)
+  x = x.mean()
+  x.backward()
+  
+  loss = x.data
+  cat = np.argmax(x_l2.data, axis=1)
+  accuracy = (cat == Y).mean()
+  
+  # SGD
+  l1.data = l1.data - lr*l1.grad
+  l2.data = l2.data - lr*l2.grad
+  
+  losses.append(loss)
+  accuracies.append(accuracy)
+  t.set_description("loss %.2f accuracy %.2f" % (loss, accuracy))
+
+# numpy forward pass
+def forward(x):
+  x = x.dot(l1.data)
+  x = np.maximum(x, 0)
+  x = x.dot(l2.data)
+  return x
+
+def numpy_eval():
+  Y_test_preds_out = forward(X_test.reshape((-1, 28*28)))
+  Y_test_preds = np.argmax(Y_test_preds_out, axis=1)
+  return (Y_test == Y_test_preds).mean()
+
+print("test set accuracy is %f" % numpy_eval())
+

tensor.py

@@ -50,6 +50,10 @@ class Tensor:
       t.grad = g
       t.backward(False)
 
+  def mean(self):
+    div = Tensor(np.array([1/self.data.size]))
+    return self.sum().mul(div)
+
 class Function:
   def apply(self, arg, *x):
     ctx = Context(arg, self, *x)
@@ -107,7 +111,7 @@ class Sum(Function):
   @staticmethod
   def forward(ctx, input):
     ctx.save_for_backward(input)
-    return np.array(input.sum())
+    return np.array([input.sum()])
 
   @staticmethod
   def backward(ctx, grad_output):