different_hidden_size.py

import argparse
import numpy as np
import matplotlib.pyplot as plt

from dataset.generator import generate_data
from model.model import FullyConnectedNeuralNetwork, ConvolutionalNeuralNetwork
from model.loss import MSE_loss, MSE_loss_backward
from utils.random import set_random_seed
from utils.metrics import accuracy
from utils.plot import show_result, show_learning_curve_and_accuracy_curve


def parse_args():
    parser = argparse.ArgumentParser(description="Train a simple neural network model.")
    parser.add_argument(
        "--dataset", type=str, default="linear", help="Dataset to use. [linear | xor]"
    )
    parser.add_argument(
        "--model",
        type=str,
        default="FullyConnectedNeuralNetwork",
        help="Model to train. [FullyConnectedNeuralNetwork | ConvolutionalNeuralNetwork]",
    )
    parser.add_argument(
        "--activation",
        type=str,
        default="sigmoid",
        help="Activation function for the model. [sigmoid | relu | tanh | none]",
    )
    parser.add_argument(
        "--epochs",
        type=int,
        default=100000,
        help="Number of epochs to train the model.",
    )
    parser.add_argument(
        "--learning_rate", type=float, default=0.1, help="Learning rate for the model."
    )
    parser.add_argument(
        "--hidden_size", type=int, default=16, help="Hidden size for the model."
    )
    parser.add_argument(
        "--optimizer",
        type=str,
        default="SGD",
        help="Optimizer for the model. [SGD | Momentum]",
    )

    args = parser.parse_args()

    return args


def train(x, y, model, epochs, learning_rate):
    loss_history = []
    accuracy_history = []
    for epoch in range(1, epochs + 1):
        y_pred = model.forward(x)

        loss = MSE_loss(y_pred, y)
        loss_grad = MSE_loss_backward(y_pred, y)
        model.backward(loss_grad)

        y_pred_binary = y_pred > 0.5
        acc = accuracy(y_pred_binary, y)

        loss_history.append(loss)
        accuracy_history.append(acc)

        if epoch % 5000 == 0:
            print(f"Epoch {epoch:7}, Loss {loss:.16f}, Accuracy {acc:.2f}")

    return loss_history, accuracy_history, y_pred, y_pred_binary, model


def test(x, y, model):
    y_pred = model.forward(x)
    y_pred_binary = y_pred > 0.5

    for i in range(x.shape[0]):
        print(
            f"iter {i+1:3} | Ground truth: {y[i][0]} | Prediction: {y_pred[i][0]:.16f} |"
        )

    print("\nTesting accuracy: ", accuracy(y_pred_binary, y))


def show_different_hidden_size(dataset, results, plot_name):
    fig, axs = plt.subplots(1, 2, figsize=(12, 5))

    # Plotting loss
    for result in results:
        loss_history = result["loss_history"]
        hidden_size = result["hidden_size"]
        axs[0].plot(
            np.arange(len(loss_history)),
            loss_history,
            label=f"Hidden size: {hidden_size}",
        )
    axs[0].set_title("Learning Curve", fontsize=10)
    axs[0].set_xlabel("Epoch")
    axs[0].set_ylabel("Loss")
    axs[0].legend()

    # Plotting accuracy
    for result in results:
        accuracy_history = result["accuracy_history"]
        hidden_size = result["hidden_size"]
        axs[1].plot(
            np.arange(len(accuracy_history)),
            accuracy_history,
            label=f"Hidden size: {hidden_size}",
        )
    axs[1].set_title("Accuracy Curve", fontsize=10)
    axs[1].set_xlabel("Epoch")
    axs[1].set_ylabel("Accuracy")
    axs[1].legend()

    if plot_name:
        plt.savefig(f"assets/{plot_name}")

    plt.close()


if __name__ == "__main__":
    args = parse_args()

    # different hidden sizes
    hidden_sizes = [16, 32, 64]

    # set random
    set_random_seed()

    # dataset
    x, y = generate_data(args.dataset)

    # training
    results = []
    for hidden_size in hidden_sizes:
        # model
        model = None
        if args.model == "FullyConnectedNeuralNetwork":
            model = FullyConnectedNeuralNetwork(
                input_size=x.shape[1],
                hidden_size=hidden_size,
                output_size=1,
                activation_type=args.activation,
                learning_rate=args.learning_rate,
                optimizer=args.optimizer,
            )
        elif args.model == "ConvolutionalNeuralNetwork":
            model = ConvolutionalNeuralNetwork()
        else:
            raise ValueError("Invalid model type.")

        loss_history, accuracy_history, y_pred, y_pred_binary, model = train(
            x, y, model, args.epochs, args.learning_rate
        )

        results.append(
            {
                "loss_history": loss_history,
                "accuracy_history": accuracy_history,
                "y_pred": y_pred,
                "y_pred_binary": y_pred_binary,
                "model": model,
                "hidden_size": hidden_size,
            }
        )

    # visualization
    show_different_hidden_size(
        args.dataset,
        results,
        f"different_hidden_size/learning_curve_and_accuracy_curve_{args.dataset}_{args.model}_{args.activation}_{args.epochs}_{args.hidden_size}.png",
    )