mckayla.py

import torch.nn as nn
import torch.optim as optim
from alias_utils.data import house
from alias_utils.model import RunnerBuilder as Builder
from alias_utils.loss import MSELoss
from alias_utils.metrics import (
    Loss as LossMetric,
    RSquared as RSquaredMetric,
    DesignMatNorm as MNormMetric,
    TestLoss as TestLossMetric,
)
from matplotlib import pyplot as plt
import matplotlib as mpl

# Shouldn't need this line anymore--it should now check and automatically disable TeX if it's not working.
mpl.rcParams["text.usetex"] = False


def train_with_optimizer(optimizer, optimizer_name):
    train_data, test_data = house.data(train_percent=0.7)
    wrapper = (
        Builder()
        .name(f"Model-{optimizer_name}")
        .loss(MSELoss())
        .optimizer(optimizer)
        .steps(
            nn.Linear(13, 8),
            nn.Linear(8, 8),
            nn.Linear(8, 1),
        )
        .with_metric(LossMetric())
        .with_metric(RSquaredMetric())
        .with_metric(MNormMetric())
        .with_metric(TestLossMetric(loss=MSELoss(), data=test_data))
        .build()
    )
    wrapper.train(train_data, n_epochs=5000)

    wrapper.plot(log=False, normalize=True)
    wrapper.plot_two(LossMetric, RSquaredMetric)
    wrapper.plot_two(LossMetric, RSquaredMetric, log=True)

    wrapper.plot_two(LossMetric, MNormMetric)
    wrapper.plot_two(LossMetric, TestLossMetric)
    plt.tight_layout()
    plt.show()

    print("Success")


def adam():
    # Read in the data
    train_data, test_data = house.data(train_percent=0.3)

    # Construct a basic layered model
    wrapper = (
        Builder()
        .name("Basic 3-Layer")
        .loss(MSELoss())
        .optimizer(optim.Adam)
        .steps(
            nn.Linear(13, 8),
            nn.Linear(8, 8),
            nn.Linear(8, 1),
        )
        .with_metric(LossMetric())
        .with_metric(RSquaredMetric())
        .with_metric(MNormMetric())
        .build()
    )

    # Train the model
    wrapper.train(train_data, n_epochs=5000)

    # print(wrapper.model.M)

    wrapper.plot(log=False, normalize=True)
    wrapper.plot_two(LossMetric, RSquaredMetric)
    wrapper.plot_two(LossMetric, RSquaredMetric, log=True)

    wrapper.plot_two(LossMetric, MNormMetric)
    wrapper.plot_two(LossMetric, MNormMetric, log=True)
    plt.tight_layout()
    plt.show()

    # test the model
    outputs = wrapper.test(test_data)
    print(len(outputs))
    y_pred, test_y = outputs[0]
    y_pred = y_pred.squeeze(dim=1)
    y_pred.detach()
    test_y.detach()
    # print('y_pred',y_pred)
    # print('test_y',test_y)

    # get the metrics
    test_loss = MSELoss()
    test_loss_vals = test_loss(y_pred=y_pred, y_true=test_y)
    print()
    print(test_loss_vals)
    print()
    plt.plot(test_loss_vals.detach().numpy())
    plt.title("Test MSE Loss")
    # plt.xlabel("")
    plt.ylabel("Loss")
    plt.show()

    print("Success")

    # # Read in the data
    # data = house.data()

    # # Construct a basic layered model
    # wrapper = (
    #     Builder()
    #     .name("Basic 3-Layer")
    #     .loss(MSELoss())
    #     .optimizer(optim.Adam)
    #     .steps(
    #         nn.Linear(13, 8),
    #         nn.Linear(8, 8),
    #         nn.Linear(8, 1),
    #     )
    #     .with_metric(LossMetric())
    #     .with_metric(RSquaredMetric())
    #     .build()
    # )

    # # Train the model
    # wrapper.train(data, n_epochs=5000)

    # # wrapper.plot(log=False, normalize=True)
    # wrapper.plot_two(LossMetric, RSquaredMetric)
    # wrapper.plot_two(LossMetric, RSquaredMetric, log=True)
    # wrapper.plot_two(0, 1)
    # plt.tight_layout()
    # plt.show()

    # print("Success")


if __name__ == "__main__":
    # adam()
    train_with_optimizer(optim.Adam, "Adam Optimizer")
    # train_with_optimizer(optim.SGD, "Stochastic Gradient Descent Optimizer")  #is not working
    # train_with_optimizer(optim.AdamW, "Adam with Weight Decay Optimizer")
    # train_with_optimizer(optim.Adagrad, "Adaptive Gradient Optimizer")
    # train_with_optimizer(optim.Adadelta, "Adadelta Optimizer")
    # train_with_optimizer(optim.ASGD, "Averaged Stochastic Gradient Descent Optimizer")
    # train_with_optimizer(optim.NAdam, "Nesterov-Accelerated Adam Optimizer")
    # train_with_optimizer(optim.LBFGS, "Limited-Memory BFGS Optimizer")