Adapting CPU script, fixing bug in slicing

Author: RobBa

examples/mnist.py

@@ -4,6 +4,7 @@
 sys.path.insert(0, str(Path(__file__).parent.parent / "python_lib"))
 
 import math
+import random
 import numpy as np
 from sklearn.datasets import fetch_openml
 from sklearn.model_selection import train_test_split
@@ -20,7 +21,7 @@
 def load_mnist():
     print("Loading MNIST...")
     mnist = fetch_openml("mnist_784", version=1, as_frame=False)
-    x = mnist.data.astype(np.float32) / 255.0  # normalize to [0,1]
+    x = mnist.data.astype(np.float32) / 255.0
     y = mnist.target.astype(np.int32)
     return x, y
 
@@ -32,16 +33,6 @@ def to_one_hot(y, n_classes=10):
     return one_hot
 
 
-def make_batches(x, y, batch_size, shuffle=True):
-    n = x.shape[0]
-    indices = np.arange(n)
-    if shuffle:
-        np.random.shuffle(indices)
-    for start in range(0, n, batch_size):
-        batch_idx = indices[start : start + batch_size]
-        yield x[batch_idx], y[batch_idx]
-
-
 # ─── network ─────────────────────────────────────────────────────────────────
 
 def make_net():
@@ -67,18 +58,24 @@ def print_weight_stats(net, batch_num):
 
 # ─── training ────────────────────────────────────────────────────────────────
 
-def train_epoch(net, loss_fn, optim, x, y, batch_size=64):
-    #print_weight_stats(net, 0)
+def train_epoch(net, loss_fn, optim, x_t, y_t, batch_size=64):
+    n = x_t.dims[0]
+
+    indices = list(range(n))
+    random.shuffle(indices)
+    x_shuf = x_t.slice(indices)
+    y_shuf = y_t.slice(indices)
 
     total_loss = 0.0
     n_batches = 0
-    max_batches = math.ceil(x.shape[0] / batch_size)
-    for xb, yb in make_batches(x, y, batch_size):
-        xTensor = fromNumpy(xb)
-        yTensor = fromNumpy(yb)
+    max_batches = math.ceil(n / batch_size)
+    for start in range(0, n, batch_size):
+        end = min(start + batch_size, n)
+        xb = x_shuf.slice(start, end)
+        yb = y_shuf.slice(start, end)
 
-        pred = net.forward(xTensor)
-        loss = loss_fn(yTensor, pred)
+        pred = net.forward(xb)
+        loss = loss_fn(yb, pred)
         loss.backward()
 
         optim.clipGradients(1.0)
@@ -89,30 +86,25 @@ def train_epoch(net, loss_fn, optim, x, y, batch_size=64):
         n_batches += 1
         if n_batches == 1 or n_batches % 10 == 0:
             print(f"Batch {n_batches} / {max_batches}, loss {loss.getitem(0)}")
-        #print_weight_stats(net, n_batches)
 
     return total_loss / n_batches
 
 
-def evaluate(net, x, y_int, batch_size=256):
+def evaluate(net, x_t, y_np, batch_size=256):
+    n = x_t.dims[0]
     correct = 0
-    total = 0
-    for xb, yb in make_batches(x, y_int, batch_size, shuffle=False):
-        xTensor = fromNumpy(xb)
-        pred = net.forward(xTensor)
+    for start in range(0, n, batch_size):
+        end = min(start + batch_size, n)
+        pred = net.forward(x_t.slice(start, end))
         pred_np = toNumpy(pred)
-
         predicted = np.argmax(pred_np, axis=1)
-        
-        correct += np.sum(predicted == np.argmax(yb, axis=1))
-        total += len(yb)
-    return correct / total
+        correct += np.sum(predicted == np.argmax(y_np[start:end], axis=1))
+    return correct / n
 
 
 # ─── main ────────────────────────────────────────────────────────────────────
 
 if __name__ == "__main__":
-    # load and split data
     x, y_int = load_mnist()
     y = to_one_hot(y_int)
 
@@ -122,16 +114,18 @@ def evaluate(net, x, y_int, batch_size=256):
 
     print(f"Train: {x_train.shape}, Val: {x_val.shape}")
 
-    # setup
+    x_train_t = fromNumpy(x_train)
+    y_train_t = fromNumpy(y_train)
+    x_val_t   = fromNumpy(x_val)
+
     net = make_net()
     loss_fn = CrossEntropyWithSoftmax()
-    optim = RmsProp(net.parameters(), 0.0001, 0.999)  # lr and decay
+    optim = RmsProp(net.parameters(), 0.0001, 0.999)
 
-    # training loop
     n_epochs = 5
     for epoch in range(n_epochs):
-        train_loss = train_epoch(net, loss_fn, optim, x_train, y_train)
-        val_acc = evaluate(net, x_val, y_val)
+        train_loss = train_epoch(net, loss_fn, optim, x_train_t, y_train_t)
+        val_acc = evaluate(net, x_val_t, y_val)
         print(
             f"Epoch {epoch+1}/{n_epochs} "
             f"loss={train_loss:.4f} "

src/backend/data_modeling/tensor.cpp

@@ -948,11 +948,12 @@ Tensor Tensor::getSlice(const tensorSize_t low, const tensorSize_t high) const {
   }
 
   makeContiguous();
-  
+
+  const tensorSize_t stride = dims.getStride(0);
   auto resDims = dims.toVector();
   resDims[0] = high-low;
   Tensor res(std::move(resDims), values->getDevice(), false);
-  values->copyValues(*res.values, low, high, 0);
+  values->copyValues(*res.values, low * stride, high * stride, 0);
   return res;
 }