update rmsprop and readme
MODIFY
README.md
README.md
@@ -41,11 +41,11 @@ print(y.grad) # dz/dy
It turns out, a decent autograd tensor library is 90% of what you need for neural networks. Add an optimizer (SGD and Adam implemented) from tinygrad.optim, write some boilerplate minibatching code, and you have all you need.
-### Neural network example (from test/mnist.py)
+### Neural network example (from test/test_mnist.py)
```python
from tinygrad.tensor import Tensor
-import tinygrad.optim as optim
+import tinygrad.optim as tinygrad_optim
from tinygrad.utils import layer_init_uniform
class TinyBobNet:
@@ -57,7 +57,8 @@ class TinyBobNet:
return x.dot(self.l1).relu().dot(self.l2).logsoftmax()
model = TinyBobNet()
-optim = optim.SGD([model.l1, model.l2], lr=0.001)
+optim = tinygrad_optim.SGD([model.l1, model.l2], lr=0.001) # or
+optim = tinygrad_optim.RMSprop([model.l1, model.l2], lr=0.001)
# ... and complete like pytorch, with (x,y) data

MODIFY
test/test_mnist.py
test/test_mnist.py
@@ -36,19 +36,8 @@ class TinyConvNet:
class TestMNIST(unittest.TestCase):
-
def test_mnist(self):
- if os.getenv("CONV") == "1":
- model = TinyConvNet()
- optim = tinygrad_optim.Adam([model.c1, model.l1, model.l2], lr=0.001)
- steps = 400
- else:
- model = TinyBobNet()
- # optim = tinygrad_optim.SGD([model.l1, model.l2], lr=0.001)
- optim = tinygrad_optim.RMSprop([model.l1, model.l2], lr=0.001)
- steps = 1000
-
- BS = 128
+ 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))
@@ -77,7 +66,7 @@ class TestMNIST(unittest.TestCase):
accuracies.append(accuracy)
t.set_description("loss %.2f accuracy %.2f" % (loss, accuracy))
- # evaluate
+ 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)
@@ -87,6 +76,24 @@ class TestMNIST(unittest.TestCase):
print("test set accuracy is %f" % accuracy)
assert accuracy > 0.95
+ # models
+ if os.getenv("CONV") == "1":
+ model = TinyConvNet()
+ optim = tinygrad_optim.Adam([model.c1, model.l1, model.l2], lr=0.001)
+ steps = 400
+ train(model, optim, steps)
+ evaluate(model)
+ else:
+ model = TinyBobNet()
+ optim = tinygrad_optim.SGD([model.l1, model.l2], lr=0.001)
+ steps = 1000
+ train(model, optim, steps)
+ evaluate(model)
+
+ # RMSprop
+ optim = tinygrad_optim.RMSprop([model.l1, model.l2], lr=0.001)
+ train(model, optim, steps)
+ evaluate(model)
if __name__ == '__main__':
unittest.main()

update rmsprop and readme