ahmet / multilayer-perceptron

import argparse
import os
import sys

import numpy as np
import pandas as pd

MALIGNANT = 1
METRICS = ("loss", "acc", "precision", "recall", "f1")


def parse_args():
    parser = argparse.ArgumentParser(description="Train a multilayer perceptron")
    parser.add_argument("--dataset", default="data_training.csv")
    parser.add_argument("--valid", default="data_validation.csv")
    parser.add_argument("--layer", type=int, nargs="+", default=[24, 24],
                        help="Hidden layer sizes (e.g. --layer 24 24 24)")
    parser.add_argument("--epochs", type=int, default=300)
    parser.add_argument("--batch_size", type=int, default=8)
    parser.add_argument("--learning_rate", type=float, default=0.1)
    parser.add_argument("--loss", default="categoricalCrossentropy",
                        choices=["categoricalCrossentropy"])
    parser.add_argument("--optimizer", default="sgd",
                        choices=["sgd", "momentum", "nesterov", "rmsprop", "adam"])
    parser.add_argument("--weights_initializer", default="heUniform",
                        choices=["heUniform", "heNormal"])
    parser.add_argument("--early_stopping", type=int, default=0, metavar="PATIENCE",
                        help="Stop after PATIENCE epochs without val_loss improvement "
                             "and restore the best weights (0 disables it)")
    parser.add_argument("--model", default="model.npz")
    parser.add_argument("--history", default="history.csv",
                        help="Where to write the per-epoch metrics")
    parser.add_argument("--compare", nargs="+", metavar="HISTORY.CSV",
                        help="Only plot the given history files together, then exit")
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()
    for name in ("epochs", "batch_size", "learning_rate"):
        if getattr(args, name) <= 0:
            parser.error(f"{name} must be positive")
    if any(n <= 0 for n in args.layer):
        parser.error("layer sizes must be positive")
    if args.early_stopping < 0:
        parser.error("early_stopping must be zero or positive")
    return args


# ---------- data ----------

def load_data(path: str, header=None) -> pd.DataFrame:
    try:
        if not os.path.isfile(path):
            raise FileNotFoundError(f"File not found: {path}")
        if not os.access(path, os.R_OK):
            raise PermissionError(f"Read permission denied: {path}")
        return pd.read_csv(path, header=header)
    except (FileNotFoundError, PermissionError,
            pd.errors.EmptyDataError, pd.errors.ParserError) as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)


def split_xy(data: pd.DataFrame):
    labels = data.iloc[:, 0]
    try:
        if not labels.isin(["M", "B"]).all():
            raise ValueError("first column must be the diagnosis label ('M' or 'B')")
        X = data.iloc[:, 1:].to_numpy(dtype=float)
    except ValueError as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)
    return X, np.eye(2)[(labels == "M").to_numpy(dtype=int)]


def standardize(X_train, X_valid):
    mean = X_train.mean(axis=0)
    std = X_train.std(axis=0)
    std[std == 0] = 1.0
    return (X_train - mean) / std, (X_valid - mean) / std, mean, std


def iterate_batches(X, Y, batch_size, rng):
    order = rng.permutation(len(X))
    for start in range(0, len(X), batch_size):
        batch = order[start:start + batch_size]
        yield X[batch], Y[batch]


# ---------- network ----------

def sigmoid(z):
    return 1.0 / (1.0 + np.exp(-z))


def softmax(z):
    z = z - z.max(axis=1, keepdims=True)
    e = np.exp(z)
    return e / e.sum(axis=1, keepdims=True)


def init_network(layers, rng, initializer="heUniform"):
    weights, biases = [], []
    for n_in, n_out in zip(layers[:-1], layers[1:]):
        if initializer == "heUniform":
            limit = np.sqrt(6.0 / n_in)
            W = rng.uniform(-limit, limit, (n_in, n_out))
        else:
            W = rng.normal(0.0, np.sqrt(2.0 / n_in), (n_in, n_out))
        weights.append(W)
        biases.append(np.zeros(n_out))
    return weights, biases


def forward(X, weights, biases):
    activations = [X]
    a = X
    for W, b in zip(weights[:-1], biases[:-1]):
        a = sigmoid(a @ W + b)
        activations.append(a)
    activations.append(softmax(a @ weights[-1] + biases[-1]))
    return activations


def backward(Y, weights, activations):
    m = Y.shape[0]
    grads_W, grads_b = [], []
    dZ = (activations[-1] - Y) / m
    for i in range(len(weights) - 1, -1, -1):
        grads_W.insert(0, activations[i].T @ dZ)
        grads_b.insert(0, dZ.sum(axis=0))
        if i > 0:
            dZ = (dZ @ weights[i].T) * activations[i] * (1 - activations[i])
    return grads_W, grads_b


# ---------- optimizers (bonus) ----------

class Optimizer:
    """Updates every parameter array in place; `params` and `grads` are aligned lists."""

    def __init__(self, name, params, lr, beta1=0.9, beta2=0.999, eps=1e-8):
        self.name, self.lr = name, lr
        self.beta1, self.beta2, self.eps = beta1, beta2, eps
        self.velocity = [np.zeros_like(p) for p in params]
        self.square = [np.zeros_like(p) for p in params]
        self.step_count = 0

    def step(self, params, grads):
        self.step_count += 1
        for i, (param, grad) in enumerate(zip(params, grads)):
            if self.name == "sgd":
                param -= self.lr * grad
                continue
            if self.name in ("momentum", "nesterov"):
                self.velocity[i] = self.beta1 * self.velocity[i] + grad
                direction = (grad + self.beta1 * self.velocity[i]
                             if self.name == "nesterov" else self.velocity[i])
                param -= self.lr * direction
                continue
            self.square[i] = (self.beta2 * self.square[i]
                              + (1 - self.beta2) * grad ** 2)
            if self.name == "rmsprop":
                param -= self.lr * grad / (np.sqrt(self.square[i]) + self.eps)
                continue
            self.velocity[i] = self.beta1 * self.velocity[i] + (1 - self.beta1) * grad
            m_hat = self.velocity[i] / (1 - self.beta1 ** self.step_count)
            v_hat = self.square[i] / (1 - self.beta2 ** self.step_count)
            param -= self.lr * m_hat / (np.sqrt(v_hat) + self.eps)


# ---------- metrics ----------

def cross_entropy(Y, A):
    A = np.clip(A, 1e-12, 1.0)
    return -np.mean(np.sum(Y * np.log(A), axis=1))


def compute_metrics(Y, A, loss_fn=cross_entropy):
    """Loss plus accuracy / precision / recall / f1 for the malignant class (bonus)."""
    true, pred = Y.argmax(axis=1), A.argmax(axis=1)
    tp = np.sum((pred == MALIGNANT) & (true == MALIGNANT))
    fp = np.sum((pred == MALIGNANT) & (true != MALIGNANT))
    fn = np.sum((pred != MALIGNANT) & (true == MALIGNANT))
    precision = tp / (tp + fp) if tp + fp else 0.0
    recall = tp / (tp + fn) if tp + fn else 0.0
    f1 = 2 * precision * recall / (precision + recall) if precision + recall else 0.0
    return {"loss": loss_fn(Y, A), "acc": np.mean(pred == true),
            "precision": precision, "recall": recall, "f1": f1}


# ---------- io / viz ----------

def save_model(path, layers, weights, biases, mean, std):
    params = {"layers": np.array(layers), "mean": mean, "std": std}
    for i, (W, b) in enumerate(zip(weights, biases)):
        params[f"W{i}"] = W
        params[f"b{i}"] = b
    np.savez(path, **params)
    print(f"> saved model to '{path}'")


def plot_histories(histories):
    """histories: {label: DataFrame}. One label -> single run, many -> comparison (bonus)."""
    import matplotlib.pyplot as plt

    _, axes = plt.subplots(1, 2, figsize=(12, 4))
    for ax, metric, name in zip(axes, ("loss", "acc"), ("loss", "accuracy")):
        for label, history in histories.items():
            line, = ax.plot(history[metric], label=f"{label} train")
            ax.plot(history[f"val_{metric}"], "--", color=line.get_color(),
                    label=f"{label} valid")
        ax.set_xlabel("epoch")
        ax.set_ylabel(name)
        ax.legend()
        ax.grid(alpha=0.3)
    plt.tight_layout()
    plt.show()


def compare_histories(paths):
    histories = {os.path.splitext(os.path.basename(p))[0]: load_data(p, header=0)
                 for p in paths}
    plot_histories(histories)


# ---------- training ----------

def evaluate(weights, biases, datasets, loss_fn):
    metrics = {}
    for prefix, X, Y in datasets:
        A = forward(X, weights, biases)[-1]
        metrics.update({f"{prefix}{k}": v
                        for k, v in compute_metrics(Y, A, loss_fn).items()})
    return metrics


def main():
    args = parse_args()
    if args.compare:
        compare_histories(args.compare)
        return

    X_train, Y_train = split_xy(load_data(args.dataset))
    X_valid, Y_valid = split_xy(load_data(args.valid))
    X_train, X_valid, mean, std = standardize(X_train, X_valid)
    datasets = (("", X_train, Y_train), ("val_", X_valid, Y_valid))

    print(f"x_train shape : {X_train.shape}")
    print(f"x_valid shape : {X_valid.shape}")

    rng = np.random.default_rng(args.seed)
    layers = [X_train.shape[1], *args.layer, Y_train.shape[1]]
    weights, biases = init_network(layers, rng, args.weights_initializer)
    params = weights + biases
    optimizer = Optimizer(args.optimizer, params, args.learning_rate)
    loss_fn = {"categoricalCrossentropy": cross_entropy}[args.loss]

    history = {}
    best = {"val_loss": np.inf, "epoch": 0, "params": None}
    width = len(str(args.epochs))

    for epoch in range(1, args.epochs + 1):
        for X_batch, Y_batch in iterate_batches(X_train, Y_train, args.batch_size, rng):
            activations = forward(X_batch, weights, biases)
            grads_W, grads_b = backward(Y_batch, weights, activations)
            optimizer.step(params, grads_W + grads_b)

        metrics = evaluate(weights, biases, datasets, loss_fn)
        for key, value in metrics.items():
            history.setdefault(key, []).append(value)
        print(f"epoch {epoch:0{width}d}/{args.epochs} - "
              + " - ".join(f"{k}: {metrics[k]:.4f}"
                           for k in ("loss", "acc", "val_loss", "val_acc")))

        if args.early_stopping:
            if metrics["val_loss"] < best["val_loss"]:
                best = {"val_loss": metrics["val_loss"], "epoch": epoch,
                        "params": [p.copy() for p in params]}
            elif epoch - best["epoch"] >= args.early_stopping:
                print(f"> early stopping: no improvement for {args.early_stopping} "
                      f"epochs, restoring epoch {best['epoch']} "
                      f"(val_loss: {best['val_loss']:.4f})")
                for param, saved in zip(params, best["params"]):
                    np.copyto(param, saved)
                break

    final = evaluate(weights, biases, datasets, loss_fn)
    print("final metrics")
    for prefix, name in (("", "train"), ("val_", "valid")):
        print(f"  {name} - "
              + " - ".join(f"{k}: {final[prefix + k]:.4f}" for k in METRICS))

    pd.DataFrame(history).to_csv(args.history, index=False)
    print(f"> saved metric history to '{args.history}'")
    save_model(args.model, layers, weights, biases, mean, std)
    plot_histories({os.path.splitext(os.path.basename(args.history))[0]:
                    pd.DataFrame(history)})


if __name__ == "__main__":
    main()