mirror of https://github.com/commaai/tinygrad.git
add mnist example
This commit is contained in:
parent
5939427795
commit
1ea9ab3e9c
|
@ -0,0 +1,78 @@
|
||||||
|
#!/usr/bin/env python
|
||||||
|
import numpy as np
|
||||||
|
from tensor import Tensor
|
||||||
|
from tqdm import trange
|
||||||
|
|
||||||
|
# load the mnist dataset
|
||||||
|
|
||||||
|
def fetch(url):
|
||||||
|
import requests, gzip, os, hashlib, numpy
|
||||||
|
fp = os.path.join("/tmp", hashlib.md5(url.encode('utf-8')).hexdigest())
|
||||||
|
if not os.path.isfile(fp):
|
||||||
|
with open(fp, "rb") as f:
|
||||||
|
dat = f.read()
|
||||||
|
else:
|
||||||
|
with open(fp, "wb") as f:
|
||||||
|
dat = requests.get(url).content
|
||||||
|
f.write(dat)
|
||||||
|
return numpy.frombuffer(gzip.decompress(dat), dtype=np.uint8).copy()
|
||||||
|
X_train = fetch("http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz")[0x10:].reshape((-1, 28, 28))
|
||||||
|
Y_train = fetch("http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz")[8:]
|
||||||
|
X_test = fetch("http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz")[0x10:].reshape((-1, 28, 28))
|
||||||
|
Y_test = fetch("http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz")[8:]
|
||||||
|
|
||||||
|
# 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)
|
||||||
|
|
||||||
|
l1 = Tensor(layer_init(784, 128))
|
||||||
|
l2 = Tensor(layer_init(128, 10))
|
||||||
|
|
||||||
|
lr = 0.01
|
||||||
|
BS = 128
|
||||||
|
losses, accuracies = [], []
|
||||||
|
for i in (t := trange(1000)):
|
||||||
|
samp = np.random.randint(0, X_train.shape[0], size=(BS))
|
||||||
|
|
||||||
|
x = Tensor(X_train[samp].reshape((-1, 28*28)))
|
||||||
|
Y = Y_train[samp]
|
||||||
|
y = np.zeros((len(samp),10), np.float32)
|
||||||
|
y[range(y.shape[0]),Y] = -1.0
|
||||||
|
y = Tensor(y)
|
||||||
|
|
||||||
|
x = x.dot(l1)
|
||||||
|
x = x.relu()
|
||||||
|
x = x_l2 = x.dot(l2)
|
||||||
|
x = x.logsoftmax()
|
||||||
|
x = x.mul(y)
|
||||||
|
x = x.mean()
|
||||||
|
x.backward()
|
||||||
|
|
||||||
|
loss = x.data
|
||||||
|
cat = np.argmax(x_l2.data, axis=1)
|
||||||
|
accuracy = (cat == Y).mean()
|
||||||
|
|
||||||
|
# SGD
|
||||||
|
l1.data = l1.data - lr*l1.grad
|
||||||
|
l2.data = l2.data - lr*l2.grad
|
||||||
|
|
||||||
|
losses.append(loss)
|
||||||
|
accuracies.append(accuracy)
|
||||||
|
t.set_description("loss %.2f accuracy %.2f" % (loss, accuracy))
|
||||||
|
|
||||||
|
# numpy forward pass
|
||||||
|
def forward(x):
|
||||||
|
x = x.dot(l1.data)
|
||||||
|
x = np.maximum(x, 0)
|
||||||
|
x = x.dot(l2.data)
|
||||||
|
return x
|
||||||
|
|
||||||
|
def numpy_eval():
|
||||||
|
Y_test_preds_out = forward(X_test.reshape((-1, 28*28)))
|
||||||
|
Y_test_preds = np.argmax(Y_test_preds_out, axis=1)
|
||||||
|
return (Y_test == Y_test_preds).mean()
|
||||||
|
|
||||||
|
print("test set accuracy is %f" % numpy_eval())
|
||||||
|
|
|
@ -50,6 +50,10 @@ class Tensor:
|
||||||
t.grad = g
|
t.grad = g
|
||||||
t.backward(False)
|
t.backward(False)
|
||||||
|
|
||||||
|
def mean(self):
|
||||||
|
div = Tensor(np.array([1/self.data.size]))
|
||||||
|
return self.sum().mul(div)
|
||||||
|
|
||||||
class Function:
|
class Function:
|
||||||
def apply(self, arg, *x):
|
def apply(self, arg, *x):
|
||||||
ctx = Context(arg, self, *x)
|
ctx = Context(arg, self, *x)
|
||||||
|
@ -107,7 +111,7 @@ class Sum(Function):
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def forward(ctx, input):
|
def forward(ctx, input):
|
||||||
ctx.save_for_backward(input)
|
ctx.save_for_backward(input)
|
||||||
return np.array(input.sum())
|
return np.array([input.sum()])
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def backward(ctx, grad_output):
|
def backward(ctx, grad_output):
|
||||||
|
|
Loading…
Reference in New Issue