refactor better
MODIFY
test/mnist.py
test/mnist.py
#!/usr/bin/env python
import numpy as np
from tinygrad.tensor import Tensor
-from tinygrad.nn import layer_init, SGD
from tinygrad.utils import fetch_mnist
+import tinygrad.optim as optim
from tqdm import trange
@@ -12,6 +12,10 @@ X_train, Y_train, X_test, Y_test = fetch_mnist()
# 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)
+
class TinyBobNet:
def __init__(self):
self.l1 = Tensor(layer_init(784, 128))
@@ -22,9 +26,8 @@ class TinyBobNet:
# optimizer
-
model = TinyBobNet()
-optim = SGD([model.l1, model.l2], lr=0.01)
+optim = optim.SGD([model.l1, model.l2], lr=0.01)
BS = 128
losses, accuracies = [], []

ADD
tinygrad/__init__.py
tinygrad/__init__.py
RENAME
tinygrad/nn.py
tinygrad/optim.py
-import numpy as np
-
-def layer_init(m, h):
- ret = np.random.uniform(-1., 1., size=(m,h))/np.sqrt(m*h)
- return ret.astype(np.float32)
-
class SGD:
def __init__(self, tensors, lr):
self.tensors = tensors

refactor better