Fix GCN validation and test evaluation

gcn/main.py

@@ -26,10 +26,23 @@ def eval_fn(x, y):
     return mx.mean(mx.argmax(x, axis=1) == y)
 
 
+def evaluate(model, x, adj, y, mask):
+    training = model.training
+    model.eval()
+    try:
+        y_hat = model(x, adj)
+        loss = loss_fn(y_hat[mask], y[mask])
+        accuracy = eval_fn(y_hat[mask], y[mask])
+        mx.eval(loss, accuracy)
+    finally:
+        model.train(training)
+    return loss, accuracy
+
+
 def forward_fn(gcn, x, adj, y, train_mask, weight_decay):
     y_hat = gcn(x, adj)
     loss = loss_fn(y_hat[train_mask], y[train_mask], weight_decay, gcn.parameters())
-    return loss, y_hat
+    return loss
 
 
 def main(args):
@@ -54,24 +67,21 @@ def main(args):
     @partial(mx.compile, inputs=state, outputs=state)
     def step():
         loss_and_grad_fn = nn.value_and_grad(gcn, forward_fn)
-        (loss, y_hat), grads = loss_and_grad_fn(
-            gcn, x, adj, y, train_mask, args.weight_decay
-        )
+        loss, grads = loss_and_grad_fn(gcn, x, adj, y, train_mask, args.weight_decay)
         optimizer.update(gcn, grads)
-        return loss, y_hat
+        return loss
 
     best_val_loss = float("inf")
     cnt = 0
 
     # Training loop
     for epoch in range(args.epochs):
         tic = time.time()
-        loss, y_hat = step()
+        loss = step()
         mx.eval(state)
 
         # Validation
-        val_loss = loss_fn(y_hat[val_mask], y[val_mask])
-        val_acc = eval_fn(y_hat[val_mask], y[val_mask])
+        val_loss, val_acc = evaluate(gcn, x, adj, y, val_mask)
         toc = time.time()
 
         # Early stopping
@@ -96,9 +106,7 @@ def step():
         )
 
     # Test
-    test_y_hat = gcn(x, adj)
-    test_loss = loss_fn(y_hat[test_mask], y[test_mask])
-    test_acc = eval_fn(y_hat[test_mask], y[test_mask])
+    test_loss, test_acc = evaluate(gcn, x, adj, y, test_mask)
 
     print(f"Test loss: {test_loss.item():.3f}  |  Test acc: {test_acc.item():.2f}")
 

gcn/test.py

@@ -0,0 +1,38 @@
+import unittest
+
+import mlx.core as mx
+import mlx.nn as nn
+from main import evaluate
+
+
+class ModeSensitiveModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.forward_modes = []
+
+    def __call__(self, x, adj):
+        self.forward_modes.append(self.training)
+        if self.training:
+            return mx.array([[0.0, 1.0], [1.0, 0.0]])
+        return mx.array([[1.0, 0.0], [0.0, 1.0]])
+
+
+class TestEvaluation(unittest.TestCase):
+    def test_evaluate_uses_eval_mode_and_restores_model_state(self):
+        labels = mx.array([0, 1])
+        mask = mx.arange(2)
+
+        for training in (True, False):
+            with self.subTest(training=training):
+                model = ModeSensitiveModel().train(training)
+
+                loss, accuracy = evaluate(model, None, None, labels, mask)
+
+                self.assertEqual(model.forward_modes, [False])
+                self.assertEqual(model.training, training)
+                self.assertAlmostEqual(accuracy.item(), 1.0)
+                self.assertAlmostEqual(loss.item(), 0.31326166, places=6)
+
+
+if __name__ == "__main__":
+    unittest.main()