github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-09 15:08:02 -05:00
Code Issues Packages Projects Releases Wiki Activity

fixup training loop

Browse Source
This commit is contained in:
George Hotz
2020-12-27 18:35:56 -05:00
parent f15bec6dbc
commit a361ef6861
4 changed files with 80 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
extra/training.py
View File

@@ -4,20 +4,25 @@ from tqdm import trange
from extra.utils import get_parameters
from tinygrad.tensor import Tensor, GPU, Device
def train(model, X_train, Y_train, optim, steps, num_classes=None, BS=128, device=Device.CPU, lossfn = lambda out,y: out.mul(y).mean()):
def sparse_categorical_crossentropy(out, Y):
num_classes = out.shape[-1]
YY = Y.flatten()
y = np.zeros((YY.shape[0], num_classes), np.float32)
# correct loss for NLL, torch NLL loss returns one per row
y[range(y.shape[0]),YY] = -1.0*num_classes
y = y.reshape(list(Y.shape)+[num_classes])
y = Tensor(y, device=out.device)
return out.mul(y).mean()
def train(model, X_train, Y_train, optim, steps, BS=128, device=Device.CPU, lossfn=sparse_categorical_crossentropy):
if device == Device.GPU: [x.gpu_() for x in get_parameters([model, optim])]
elif device == Device.ANE: [x.ane_() for x in get_parameters([model, optim])]
if num_classes is None: num_classes = Y_train.max().astype(int)+1
losses, accuracies = [], []
for i in (t := trange(steps, disable=os.getenv('CI') is not None)):
samp = np.random.randint(0, X_train.shape[0], size=(BS))
x = Tensor(X_train[samp].reshape((-1, 28*28)).astype(np.float32), device=device)
Y = Y_train[samp]
y = np.zeros((len(samp),num_classes), np.float32)
# correct loss for NLL, torch NLL loss returns one per row
y[range(y.shape[0]),Y] = -1.0*num_classes
y = Tensor(y, device=device)
x = Tensor(X_train[samp], device=device)
y = Y_train[samp]
# network
out = model.forward(x)
@@ -29,7 +34,7 @@ def train(model, X_train, Y_train, optim, steps, num_classes=None, BS=128, devic
optim.step()
cat = np.argmax(out.cpu().data, axis=1)
accuracy = (cat == Y).mean()
accuracy = (cat == y).mean()
# printing
loss = loss.cpu().data
@@ -41,7 +46,7 @@ def evaluate(model, X_test, Y_test, num_classes=None, device=Device.CPU, BS=128)
def numpy_eval(num_classes):
Y_test_preds_out = np.zeros((len(Y_test),num_classes))
for i in trange(len(Y_test)//BS, disable=os.getenv('CI') is not None):
Y_test_preds_out[i*BS:(i+1)*BS] = model.forward(Tensor(X_test[i*BS:(i+1)*BS].reshape((-1, 28*28)).astype(np.float32), device=device)).cpu().data
Y_test_preds_out[i*BS:(i+1)*BS] = model.forward(Tensor(X_test[i*BS:(i+1)*BS], device=device)).cpu().data
Y_test_preds = np.argmax(Y_test_preds_out, axis=1)
return (Y_test == Y_test_preds).mean()