calvin.py

import torch
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, MNormLoss
from alias_utils.metrics import (
    Loss as LossMetric,
    # RSquared as RSquaredMetric,
    DesignMatNorm as MNormMetric,
)
from matplotlib import pyplot as plt

if __name__ == "__main__":
    # Read in the data
    data = house.data()

    torch.manual_seed(13)
    # 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(name="Standard Loss"))
        .with_metric(MNormMetric(name="Standard MNorm"))
        .build()
    )

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

    torch.manual_seed(13)
    # Construct model w m norm loss
    m_wrapper = (
        Builder()
        .name("Basic 3-Layer")
        .loss(MNormLoss(alpha=1e4))
        .optimizer(optim.Adam)
        .steps(
            nn.Linear(13, 8),
            nn.Linear(8, 8),
            nn.Linear(8, 1),
        )
        .with_metric(LossMetric(name="MNorm Loss"))
        .with_metric(MNormMetric(name="MNorm MNorm"))
        .build()
    )

    # Train the model
    m_wrapper.train(data, n_epochs=1000)

    # plot loss and m norm for each
    # wrapper.plot_two(LossMetric, MNormMetric, log=True)
    # m_wrapper.plot_two(LossMetric, MNormMetric, log=True)
    # plt.tight_layout()
    # plt.show()

    # plot both on same plot
    plt.title("Comparison w and w/o MNorm")
    wrapper.plot(log=True, normalize=False)
    m_wrapper.plot(log=True, normalize=False)
    plt.show()

    print("Success")