clean up test

test/test_mnist.py

@@ -6,7 +6,6 @@ from tinygrad.tensor import Tensor
 from tinygrad.utils import layer_init_uniform, fetch_mnist
 import tinygrad.optim as optim
 from tqdm import trange
-np.random.seed(1337)
 
 # load the mnist dataset
 X_train, Y_train, X_test, Y_test = fetch_mnist()
@@ -21,7 +20,6 @@ class TinyBobNet:
     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):
     conv = 7
@@ -36,65 +34,67 @@ class TinyConvNet:
     x = x.reshape(Tensor(np.array((x.shape[0], -1))))
     return x.dot(self.l1).relu().dot(self.l2).logsoftmax()
 
+def train(model, optim, steps, BS=128):
+  losses, accuracies = [], []
+  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)).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
+    y[range(y.shape[0]),Y] = -10.0
+    y = Tensor(y)
+    
+    # network
+    out = model.forward(x)
+
+    # NLL loss function
+    loss = out.mul(y).mean()
+    loss.backward()
+    optim.step()
+    
+    cat = np.argmax(out.data, axis=1)
+    accuracy = (cat == Y).mean()
+    
+    # printing
+    loss = loss.data
+    losses.append(loss)
+    accuracies.append(accuracy)
+    t.set_description("loss %.2f accuracy %.2f" % (loss, accuracy))
+
+def evaluate(model):
+  def numpy_eval():
+    Y_test_preds_out = model.forward(Tensor(X_test.reshape((-1, 28*28)).astype(np.float32)))
+    Y_test_preds = np.argmax(Y_test_preds_out.data, axis=1)
+    return (Y_test == Y_test_preds).mean()
+
+  accuracy = numpy_eval()
+  print("test set accuracy is %f" % accuracy)
+  assert accuracy > 0.95
 
 class TestMNIST(unittest.TestCase):
-  def test_mnist(self):
-    def train(model, optim, steps, BS=128):
-      losses, accuracies = [], []
-      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)).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
-        y[range(y.shape[0]),Y] = -10.0
-        y = Tensor(y)
-        
-        # network
-        out = model.forward(x)
-
-        # NLL loss function
-        loss = out.mul(y).mean()
-        loss.backward()
-        optim.step()
-        
-        cat = np.argmax(out.data, axis=1)
-        accuracy = (cat == Y).mean()
-        
-        # printing
-        loss = loss.data
-        losses.append(loss)
-        accuracies.append(accuracy)
-        t.set_description("loss %.2f accuracy %.2f" % (loss, accuracy))
-
-    def evaluate(model):
-      def numpy_eval():
-        Y_test_preds_out = model.forward(Tensor(X_test.reshape((-1, 28*28)).astype(np.float32)))
-        Y_test_preds = np.argmax(Y_test_preds_out.data, axis=1)
-        return (Y_test == Y_test_preds).mean()
-
-      accuracy = numpy_eval()
-      print("test set accuracy is %f" % accuracy)
-      assert accuracy > 0.95
-
-    # models
+  def test_mnist_conv(self):
+    np.random.seed(1337)
     model = TinyConvNet()
     optimizer = optim.Adam([model.c1, model.l1, model.l2], lr=0.001)
-    steps = 400
-    train(model, optimizer, steps)
+    train(model, optimizer, steps=400)
     evaluate(model)
     
+  def test_mnist_sgd(self):
+    np.random.seed(1337)
     model = TinyBobNet()
-    steps = 1000
     optimizer = optim.SGD([model.l1, model.l2], lr=0.001)
-    train(model, optimizer, steps)
+    train(model, optimizer, steps=1000)
     evaluate(model)
     
+  def test_mnist_rmsprop(self):
+    np.random.seed(1337)
     model = TinyBobNet()
     optimizer = optim.RMSprop([model.l1, model.l2], lr=0.001)
-    train(model, optimizer, steps)
+    train(model, optimizer, steps=1000)
     evaluate(model)
 
 if __name__ == '__main__':
   unittest.main()
+