stunning_mnist [run_process_replay] (#6828)

* stunning_mnist [run_process_replay]

* add loss to stunning mnist

examples/stunning_mnist.py

@@ -0,0 +1,53 @@
+# beautiful mnist in the new "one-shot" style
+# one realize in the whole graph
+# depends on:
+#  - "big graph" UOp scheduling
+#  - symbolic removal
+
+from examples.beautiful_mnist import Model
+from tinygrad import Tensor, nn, getenv, GlobalCounters
+from tinygrad.nn.datasets import mnist
+from tinygrad.helpers import trange, DEBUG
+
+if __name__ == "__main__":
+  X_train, Y_train, X_test, Y_test = mnist()
+  print("*** got data")
+
+  model = Model()
+  print("*** got model")
+
+  opt = nn.optim.Adam(nn.state.get_parameters(model))
+  print("*** got optimizer")
+
+  samples = Tensor.randint(getenv("STEPS", 10), getenv("BS", 512), high=X_train.shape[0])
+  X_samp, Y_samp = X_train[samples], Y_train[samples]
+  print("*** got samples")
+
+  with Tensor.train():
+    # TODO: this shouldn't be a for loop. something like: (contract is still up in the air)
+    """
+    i = UOp.range(samples.shape[0])  # TODO: fix range function on UOp
+    losses = model(X_samp[i]).sparse_categorical_crossentropy(Y_samp[i]).backward().contract(i)
+    opt.schedule_steps(i)
+    """
+    losses = []
+    for i in range(samples.shape[0]):
+      opt.zero_grad()
+      losses.append(model(X_samp[i]).sparse_categorical_crossentropy(Y_samp[i]).backward())
+      opt.schedule_step()
+    # TODO: this stack currently breaks the "generator" aspect of losses. it probably shouldn't
+    #losses = Tensor.stack(*losses)
+  print("*** scheduled training")
+
+  # evaluate the model
+  with Tensor.test():
+    test_acc = ((model(X_test).argmax(axis=1) == Y_test).mean()*100)
+  print("*** scheduled eval")
+
+  # NOTE: there's no kernels run in the scheduling phase
+  assert GlobalCounters.kernel_count == 0, "kernels were run during scheduling!"
+
+  # only actually do anything at the end
+  if getenv("LOSS", 1):
+    for i in (t:=trange(len(losses))): t.set_description(f"loss: {losses[i].item():6.2f}")
+  print(f"test_accuracy: {test_acc.item():5.2f}%")

test/test_schedule.py

@@ -246,7 +246,7 @@ class TestSchedule(unittest.TestCase):
 
   def test_fold_conv_batchnorm_optim(self):
     # this is too high
-    for optim, cnt in [(nn.optim.Adam, 17), (nn.optim.SGD, 15)]:
+    for optim, cnt in [(nn.optim.Adam, 18), (nn.optim.SGD, 15)]:
       with self.subTest(optim=optim.__name__):
         with Tensor.train():
           img = Tensor.ones(1,3,4,4)
@@ -913,7 +913,7 @@ class TestSchedule(unittest.TestCase):
       _realize_weights(layer)
       opt = nn.optim.Adam(nn.state.get_parameters(layer), lr=1e-4)
       layer(x).relu().sum().backward()
-      check_schedule(opt.schedule_step(), 9)
+      check_schedule(opt.schedule_step(), 10)
 
   def test_adam_conv_fuse(self):
     with Tensor.train():
@@ -923,7 +923,7 @@ class TestSchedule(unittest.TestCase):
       opt = nn.optim.Adam(nn.state.get_parameters(c1), lr=1e-4)
       opt.zero_grad()
       c1(img).relu().sum().backward()
-      check_schedule(opt.schedule_step(), 9)
+      check_schedule(opt.schedule_step(), 10)
 
   def test_adam_2convs_fuse(self):
     with Tensor.train():
@@ -934,7 +934,7 @@ class TestSchedule(unittest.TestCase):
       opt = nn.optim.Adam(nn.state.get_parameters([c1, c2]), lr=1e-4)
       opt.zero_grad()
       c2(c1(img).relu()).relu().sum().backward()
-      check_schedule(opt.schedule_step(), 12)
+      check_schedule(opt.schedule_step(), 13)
 
   def test_sgd_conv_fuse(self):
     with Tensor.train():

tinygrad/nn/optim.py

@@ -126,7 +126,7 @@ class LAMB(Optimizer):
   def __init__(self, params: List[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-6, weight_decay=0.0, adam=False):
     super().__init__(params, lr)
     self.b1, self.b2, self.eps, self.wd, self.adam = b1, b2, eps, weight_decay, adam
-    self.b1_t, self.b2_t = (Tensor([1], dtype=dtypes.float32, device=self.device, requires_grad=False).realize() for _ in [b1, b2])
+    self.b1_t, self.b2_t = (Tensor.ones((1,), dtype=dtypes.float32, device=self.device, requires_grad=False).contiguous() for _ in [b1, b2])
     self.m = [Tensor.zeros(*t.shape, dtype=dtypes.float32, device=t.device, requires_grad=False).contiguous() for t in self.params]
     self.v = [Tensor.zeros(*t.shape, dtype=dtypes.float32, device=t.device, requires_grad=False).contiguous() for t in self.params]