wip

Author: balancap

experiments/mnist/mnist_classifier_from_scratch.py

@@ -46,6 +46,7 @@ def predict(params, inputs):
     activations = inputs
     for w, b in params[:-1]:
         jsa.ops.debug_callback(partial(print_mean_std, "W"), w)
+        jsa.ops.debug_callback(partial(print_mean_std, "B"), b)
         (w,) = jsa.ops.debug_callback_grad(partial(print_mean_std, "WGrad"), w)
 
         # Matmul + relu
@@ -107,7 +108,7 @@ def update(params, batch):
         grads = grad(loss)(params, batch)
         return [(w - step_size * dw, b - step_size * db) for (w, b), (dw, db) in zip(params, grads)]
 
-    num_batches = 2
+    num_batches = 4
     num_epochs = 2
     for epoch in range(num_epochs):
         print("EPOCH:", epoch)