IQ.Pilot Release Commit @ bec7652

artifacts/package_sources/tinygrad/extra/training.py

@@ -0,0 +1,57 @@
+import numpy as np
+from tinygrad.tensor import Tensor
+from tinygrad.helpers import trange, Context
+from tinygrad.engine.jit import TinyJit
+
+
+def train(model, X_train, Y_train, optim, steps, BS=128, lossfn=lambda out,y: out.sparse_categorical_crossentropy(y),
+        transform=lambda x: x, target_transform=lambda x: x, noloss=False, allow_jit=True):
+
+  def train_step(x, y):
+    # network
+    out = model.forward(x) if hasattr(model, 'forward') else model(x)
+    loss = lossfn(out, y)
+    optim.zero_grad()
+    loss.backward()
+    if noloss: del loss
+    optim.step()
+    if noloss: return (None, None)
+    cat = out.argmax(axis=-1)
+    accuracy = (cat == y).mean()
+    return loss.realize(), accuracy.realize()
+
+  if allow_jit: train_step = TinyJit(train_step)
+
+  with Context(TRAINING=1):
+    losses, accuracies = [], []
+    for i in (t := trange(steps, disable=None)):
+      samp = np.random.randint(0, X_train.shape[0], size=(BS))
+      x = Tensor(transform(X_train[samp]))
+      y = Tensor(target_transform(Y_train[samp]))
+      loss, accuracy = train_step(x, y)
+      # printing
+      if not noloss:
+        loss, accuracy = loss.numpy(), accuracy.numpy()
+        losses.append(loss)
+        accuracies.append(accuracy)
+        t.set_description("loss %.2f accuracy %.2f" % (loss, accuracy))
+  return [losses, accuracies]
+
+
+def evaluate(model, X_test, Y_test, num_classes=None, BS=128, return_predict=False, transform=lambda x: x,
+             target_transform=lambda y: y):
+  @Context(TRAINING=0)
+  def numpy_eval(Y_test, num_classes):
+    Y_test_preds_out = np.zeros(list(Y_test.shape)+[num_classes])
+    for i in trange((len(Y_test)-1)//BS+1, disable=None):
+      x = Tensor(transform(X_test[i*BS:(i+1)*BS]))
+      out = model.forward(x) if hasattr(model, 'forward') else model(x)
+      Y_test_preds_out[i*BS:(i+1)*BS] = out.numpy()
+    Y_test_preds = np.argmax(Y_test_preds_out, axis=-1)
+    Y_test = target_transform(Y_test)
+    return (Y_test == Y_test_preds).mean(), Y_test_preds
+
+  if num_classes is None: num_classes = Y_test.max().astype(int)+1
+  acc, Y_test_pred = numpy_eval(Y_test, num_classes)
+  print("test set accuracy is %f" % acc)
+  return (acc, Y_test_pred) if return_predict else acc