ai-lab/m01_autograd/train.py

"""AI01 lab, step 8: teach a tiny network XOR with gradient descent.

Run it when step 8 passes:
    python train.py            # trains and prints the loss every 25 steps
    python train.py --viz      # also writes out/xor.html (loss curve + decision map)
"""
import argparse
import random

from nn import MLP

# XOR: the output is +1 when exactly one input is 1, and -1 otherwise.
# No single straight line separates the +1 points from the -1 points, which
# is why the network needs a hidden layer.
XS = [[0.0, 0.0], [0.0, 1.0], [1.0, 0.0], [1.0, 1.0]]
YS = [-1.0, 1.0, 1.0, -1.0]


def mse_loss(preds, targets):
    """Mean squared error: the average of (pred - target)**2 over the examples.

    preds    a list of Values (the model's outputs).
    targets  a list of floats, the same length.
    Return one Value, so you can call backward() on it.
    """
    raise NotImplementedError("Step 8: mse_loss")


def train_step(model, xs, ys, lr):
    """One step of gradient descent on the whole dataset. Return the loss as a float.

    1. Zero every parameter's grad. Never assume they are already clean.
    2. Forward: predict every example in xs and compute mse_loss against ys.
    3. Backward: loss.backward().
    4. Update: move every parameter a step of size lr against its gradient.
    Return the loss from step 2 (computed before the update), as a float.
    """
    raise NotImplementedError("Step 8: train_step")


def train(model, xs, ys, steps, lr):
    """Run train_step `steps` times. Return the list of losses, one per step."""
    raise NotImplementedError("Step 8: train")


def main():
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--steps", type=int, default=400)
    parser.add_argument("--lr", type=float, default=0.1)
    parser.add_argument("--seed", type=int, default=1)
    parser.add_argument("--viz", action="store_true", help="write out/xor.html")
    args = parser.parse_args()

    random.seed(args.seed)
    model = MLP(2, [4, 4, 1])
    losses = train(model, XS, YS, args.steps, args.lr)
    for step in range(0, args.steps, 25):
        print(f"step {step:4d}   loss {losses[step]:.5f}")
    print(f"final      loss {losses[-1]:.5f}")
    for x, y in zip(XS, YS):
        print(f"input {x}  target {y:+.0f}  prediction {model(x).data:+.3f}")

    if args.viz:
        import sys
        from pathlib import Path
        sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "tools"))
        import viz
        path = viz.plot_training(losses, model, XS, YS, "out/xor.html")
        print("wrote", path)


if __name__ == "__main__":
    main()