From 547733e57cfec2afc59d425f1f0d3f34bc9563b7 Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Tue, 1 Oct 2024 15:00:48 +0800 Subject: [PATCH] stunning_mnist [run_process_replay] (#6828) * stunning_mnist [run_process_replay] * add loss to stunning mnist --- examples/stunning_mnist.py | 53 ++++++++++++++++++++++++++++++++++++++ test/test_schedule.py | 8 +++--- tinygrad/nn/optim.py | 2 +- 3 files changed, 58 insertions(+), 5 deletions(-) create mode 100644 examples/stunning_mnist.py diff --git a/examples/stunning_mnist.py b/examples/stunning_mnist.py new file mode 100644 index 0000000000..6e9e4dc645 --- /dev/null +++ b/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}%") diff --git a/test/test_schedule.py b/test/test_schedule.py index ac809c157a..24047042c5 100644 --- a/test/test_schedule.py +++ b/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(): diff --git a/tinygrad/nn/optim.py b/tinygrad/nn/optim.py index fd9c612bcf..ff77952d70 100644 --- a/tinygrad/nn/optim.py +++ b/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]