neuralnetwork/model.py

import numpy as np


def relu(x):
    return np.maximum(0.0, x)


def relu_deriv(x):
    return (x > 0.0).astype(float)


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


def cross_entropy_loss(Y_hat, Y):
    m = Y.shape[0]
    safe_hat = np.clip(Y_hat, 1e-12, 1.0 - 1e-12)
    return -np.sum(Y * np.log(safe_hat)) / m


def he_initialization(fan_in, fan_out):
    return np.random.randn(fan_in, fan_out) * np.sqrt(2.0 / max(1, fan_in))


def xavier_initialization(fan_in, fan_out):
    return np.random.randn(fan_in, fan_out) * np.sqrt(2.0 / (fan_in + fan_out))


class NeuralNetwork:
    def __init__(
        self,
        layer_sizes,
        lr=0.01,
        optimizer="adam",
        reg_lambda=1e-4,
        dropout=0.0,
        seed=42,
    ):
        if len(layer_sizes) < 3:
            raise ValueError("layer_sizes must contain input, hidden, and output sizes")

        self.layer_sizes = list(layer_sizes)
        self.lr = float(lr)
        self.optimizer = optimizer.lower()
        self.reg_lambda = float(reg_lambda)
        self.dropout = float(dropout)

        if self.lr <= 0.0:
            raise ValueError("lr must be positive")
        if self.reg_lambda < 0.0:
            raise ValueError("reg_lambda must be non-negative")
        if not 0.0 <= self.dropout < 1.0:
            raise ValueError("dropout must be between 0 and 1")

        self.rng = np.random.RandomState(seed)
        self.params = {}
        self.cache = {}
        self._initialize_parameters()
        self._initialize_optimizer_state()

    @property
    def num_layers(self):
        return len(self.layer_sizes) - 1

    def _initialize_parameters(self):
        for layer in range(1, self.num_layers + 1):
            fan_in = self.layer_sizes[layer - 1]
            fan_out = self.layer_sizes[layer]
            if layer == self.num_layers:
                self.params[f"W{layer}"] = xavier_initialization(fan_in, fan_out)
            else:
                self.params[f"W{layer}"] = he_initialization(fan_in, fan_out)
            self.params[f"b{layer}"] = np.zeros((1, fan_out), dtype=float)

    def _initialize_optimizer_state(self):
        if self.optimizer == "adam":
            self.optimizer_state = {
                "t": 0,
                "m": {
                    **{f"W{layer}": np.zeros_like(self.params[f"W{layer}"]) for layer in range(1, self.num_layers + 1)},
                    **{f"b{layer}": np.zeros_like(self.params[f"b{layer}"]) for layer in range(1, self.num_layers + 1)},
                },
                "v": {
                    **{f"W{layer}": np.zeros_like(self.params[f"W{layer}"]) for layer in range(1, self.num_layers + 1)},
                    **{f"b{layer}": np.zeros_like(self.params[f"b{layer}"]) for layer in range(1, self.num_layers + 1)},
                },
                "beta1": 0.9,
                "beta2": 0.999,
                "epsilon": 1e-8,
            }
        elif self.optimizer == "momentum":
            self.optimizer_state = {
                "velocity": {
                    **{f"W{layer}": np.zeros_like(self.params[f"W{layer}"]) for layer in range(1, self.num_layers + 1)},
                    **{f"b{layer}": np.zeros_like(self.params[f"b{layer}"]) for layer in range(1, self.num_layers + 1)},
                },
                "beta": 0.9,
            }
        else:
            self.optimizer_state = {}

    def forward(self, X, training=False):
        self.cache["A0"] = X
        self.cache["dropout_masks"] = {}
        A = X

        for layer in range(1, self.num_layers + 1):
            Z = A.dot(self.params[f"W{layer}"]) + self.params[f"b{layer}"]
            self.cache[f"Z{layer}"] = Z

            if layer < self.num_layers:
                A = relu(Z)
                if training and self.dropout > 0.0:
                    mask = (self.rng.rand(*A.shape) >= self.dropout).astype(float) / (1.0 - self.dropout)
                    A *= mask
                    self.cache["dropout_masks"][f"D{layer}"] = mask
                self.cache[f"A{layer}"] = A
            else:
                A = softmax(Z)
                self.cache[f"A{layer}"] = A

        return A

    def compute_loss(self, Y_hat, Y):
        loss = cross_entropy_loss(Y_hat, Y)
        if self.reg_lambda > 0.0:
            reg_loss = 0.0
            for layer in range(1, self.num_layers + 1):
                reg_loss += np.sum(self.params[f"W{layer}"] ** 2)
            loss += self.reg_lambda * reg_loss / Y.shape[0]
        return loss

    def backward(self, Y):
        m = Y.shape[0]
        grads = {}
        dA = self.cache[f"A{self.num_layers}"] - Y

        for layer in reversed(range(1, self.num_layers + 1)):
            A_prev = self.cache[f"A{layer - 1}"]
            grads[f"dW{layer}"] = A_prev.T.dot(dA) / m + (self.reg_lambda / m) * self.params[f"W{layer}"]
            grads[f"db{layer}"] = np.sum(dA, axis=0, keepdims=True) / m

            if layer > 1:
                dA = dA.dot(self.params[f"W{layer}"].T)
                dZ = dA * relu_deriv(self.cache[f"Z{layer - 1}"])
                if self.dropout > 0.0:
                    dZ *= self.cache["dropout_masks"][f"D{layer - 1}"]
                dA = dZ

        self._update_parameters(grads)

    def _update_parameters(self, grads):
        if self.optimizer == "adam":
            state = self.optimizer_state
            state["t"] += 1
            beta1 = state["beta1"]
            beta2 = state["beta2"]
            eps = state["epsilon"]
            t = state["t"]

            for layer in range(1, self.num_layers + 1):
                for param_name in [f"W{layer}", f"b{layer}"]:
                    g = grads[f"d{param_name}"]
                    state["m"][param_name] = beta1 * state["m"][param_name] + (1 - beta1) * g
                    state["v"][param_name] = beta2 * state["v"][param_name] + (1 - beta2) * (g * g)
                    m_hat = state["m"][param_name] / (1 - beta1 ** t)
                    v_hat = state["v"][param_name] / (1 - beta2 ** t)
                    self.params[param_name] -= self.lr * m_hat / (np.sqrt(v_hat) + eps)
        elif self.optimizer == "momentum":
            state = self.optimizer_state
            beta = state["beta"]
            for layer in range(1, self.num_layers + 1):
                for param_name in [f"W{layer}", f"b{layer}"]:
                    velocity = state["velocity"][param_name]
                    velocity *= beta
                    velocity += self.lr * grads[f"d{param_name}"]
                    self.params[param_name] -= velocity
                    state["velocity"][param_name] = velocity
        else:
            for layer in range(1, self.num_layers + 1):
                self.params[f"W{layer}"] -= self.lr * grads[f"dW{layer}"]
                self.params[f"b{layer}"] -= self.lr * grads[f"db{layer}"]

    def fit(
        self,
        X,
        Y,
        X_val=None,
        Y_val=None,
        epochs=100,
        batch_size=16,
        patience=20,
        verbose=True,
    ): 
        n = X.shape[0]
        best_val_loss = float("inf")
        best_params = None
        no_improve = 0

        for epoch in range(1, epochs + 1):
            perm = self.rng.permutation(n)
            X_shuffled = X[perm]
            Y_shuffled = Y[perm]

            for start in range(0, n, batch_size):
                xb = X_shuffled[start : start + batch_size]
                yb = Y_shuffled[start : start + batch_size]
                self.forward(xb, training=True)
                self.backward(yb)

            Y_hat = self.forward(X)
            train_loss = self.compute_loss(Y_hat, Y)
            train_acc = np.mean(np.argmax(Y_hat, axis=1) == np.argmax(Y, axis=1))

            val_loss = None
            val_acc = None
            if X_val is not None and Y_val is not None and X_val.shape[0] > 0:
                Y_val_hat = self.forward(X_val)
                val_loss = self.compute_loss(Y_val_hat, Y_val)
                val_acc = np.mean(np.argmax(Y_val_hat, axis=1) == np.argmax(Y_val, axis=1))

                if val_loss + 1e-8 < best_val_loss:
                    best_val_loss = val_loss
                    best_params = {k: v.copy() for k, v in self.params.items()}
                    no_improve = 0
                else:
                    no_improve += 1

            if verbose and (epoch == 1 or epoch % max(1, epochs // 10) == 0):
                msg = f"Epoch {epoch}/{epochs} - train_loss: {train_loss:.4f}, train_acc: {train_acc:.4f}"
                if val_loss is not None:
                    msg += f", val_loss: {val_loss:.4f}, val_acc: {val_acc:.4f}"
                print(msg)

            if no_improve >= patience:
                if best_params is not None:
                    self.params = best_params
                break

        return self

    def predict_proba(self, X):
        return self.forward(X, training=False)

    def predict(self, X):
        return np.argmax(self.predict_proba(X), axis=1)