ai-lab/m01_autograd/tests/test_train.py

"""AI01 lab tests for train.py (step 8)."""
import random
import unittest

from _impl import load

engine = load("engine")
nn = load("nn")
train = load("train")
Value = engine.Value

XOR_X = [[0.0, 0.0], [0.0, 1.0], [1.0, 0.0], [1.0, 1.0]]
XOR_Y = [-1.0, 1.0, 1.0, -1.0]


class Step8Training(unittest.TestCase):
    def test_mse_is_the_mean_of_squared_errors(self):
        loss = train.mse_loss([Value(1.0), Value(3.0)], [0.0, 1.0])
        self.assertIsInstance(loss, Value)
        self.assertAlmostEqual(loss.data, 2.5)      # (1 + 4) / 2

    def test_mse_gradient(self):
        p0, p1 = Value(1.0), Value(3.0)
        train.mse_loss([p0, p1], [0.0, 1.0]).backward()
        self.assertAlmostEqual(p0.grad, 1.0)        # 2 (p - t) / n
        self.assertAlmostEqual(p1.grad, 2.0)

    def test_train_step_zeroes_stale_grads_and_steps_downhill(self):
        # One linear neuron, y = w x + b, starting at w = 0.5, b = 0.
        model = nn.MLP(1, [1])
        w, b = model.parameters()
        w.data, b.data = 0.5, 0.0
        w.grad, b.grad = 100.0, -100.0              # stale junk from "earlier"
        loss = train.train_step(model, [[1.0], [2.0]], [1.0, 3.0], lr=0.1)
        # preds 0.5, 1.0; loss = (0.25 + 4) / 2 = 2.125
        # dL/dw = (2/2)(-0.5*1 + -2*2) = -4.5 ; dL/db = (2/2)(-0.5 + -2) = -2.5
        self.assertIsInstance(loss, float)
        self.assertAlmostEqual(loss, 2.125)
        self.assertAlmostEqual(w.data, 0.95)        # 0.5 - 0.1 * (-4.5)
        self.assertAlmostEqual(b.data, 0.25)        # 0.0 - 0.1 * (-2.5)

    def test_train_returns_one_loss_per_step(self):
        random.seed(3)
        losses = train.train(nn.MLP(2, [3, 1]), XOR_X, XOR_Y, steps=5, lr=0.05)
        self.assertEqual(len(losses), 5)
        self.assertTrue(all(isinstance(v, float) for v in losses))

    def test_learns_xor(self):
        for seed in (0, 1, 2):
            with self.subTest(seed=seed):
                random.seed(seed)
                model = nn.MLP(2, [4, 4, 1])
                losses = train.train(model, XOR_X, XOR_Y, steps=400, lr=0.1)
                self.assertLess(losses[-1], 0.05)
                self.assertLess(losses[-1], losses[0] / 10)
                for x, y in zip(XOR_X, XOR_Y):
                    self.assertGreater(model(x).data * y, 0.5,
                                       f"prediction for {x} has the wrong sign or is too timid")


if __name__ == "__main__":
    unittest.main()