train.py

import argparse
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
from torchvision import transforms
from tqdm import tqdm

from dataset import MultiMNIST
from model import MTLModel
from pcgrad import PCGrad
from utils import AverageMeter, setup_experiment


def get_summary(meters: dict) -> str:
    return " ".join([f"{k}: {v.avg:5.7f}" for k, v in meters.items()])


def accuracy(y_pred, y_true):
    return ((y_pred.argmax(dim=-1) == y_true).float()).mean() * 100


def to_dev(batch, device) -> None:
    return [x.to(device) for x in batch]


def main(args: argparse.Namespace) -> None:
    # Set Seeds
    setup_experiment(args.seed)

    # Check CUDA
    if torch.cuda.is_available():
        device = torch.device("cuda")
    else:
        device = torch.device("cpu")

    print(f"Device: {device}")

    global_transformer = transforms.Compose(
        [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
    )

    train_dst = MultiMNIST(
        root=args.path,
        train=True,
        download=args.download,
        transform=global_transformer,
        multi=True,
    )

    train_loader = torch.utils.data.DataLoader(
        train_dst, batch_size=args.batch_size, shuffle=True, num_workers=2
    )

    val_dst = MultiMNIST(
        root=args.path,
        train=False,
        download=args.download,
        transform=global_transformer,
        multi=True,
    )

    val_loader = torch.utils.data.DataLoader(
        val_dst, batch_size=args.batch_size, shuffle=False, num_workers=2
    )

    model = MTLModel()
    model.to(device)

    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)

    if args.pcgrad:
        optimizer = PCGrad(optimizer)

    meters = {
        "t_loss_right": AverageMeter("t_loss_right", ":.4e"),
        "t_loss_left": AverageMeter("t_loss_left", ":.4e"),
        "v_loss": AverageMeter("v_loss", ":.4e"),
        "v_loss_right": AverageMeter("v_loss_right", ":.4e"),
        "v_loss_left": AverageMeter("v_loss_left", ":.4e"),
        "v_acc_right": AverageMeter("v_acc_right", ":.4e"),
        "v_acc_left": AverageMeter("v_acc_left", ":.4e"),
    }

    best_val_loss = np.inf

    for i_epoch in range(args.n_epochs):
        # Training
        model.train()

        for meter in meters.values():
            meter.reset()

        for batch in tqdm(train_loader):
            img, label_l, label_r = (x.to(device) for x in batch)

            optimizer.zero_grad()

            out_l, out_r = model(img)

            left_loss = F.nll_loss(out_l, label_l)
            right_loss = F.nll_loss(out_r, label_r)

            if args.pcgrad:
                optimizer.pc_backward(model, [left_loss, right_loss])
            else:
                (left_loss + right_loss).backward()

            optimizer.step()

            meters["t_loss_right"].update(right_loss.item())
            meters["t_loss_left"].update(left_loss.item())

        # Validation
        model.eval()

        with torch.inference_mode():
            for batch in tqdm(val_loader):
                img, label_l, label_r = (x.to(device) for x in batch)

                out_l, out_r = model(img)

                left_loss = F.nll_loss(out_l, label_l)
                right_loss = F.nll_loss(out_r, label_r)
                loss = left_loss + right_loss

                left_acc = accuracy(out_l, label_l)
                right_acc = accuracy(out_r, label_r)

                meters["v_loss_right"].update(right_loss.item())
                meters["v_loss_left"].update(left_loss.item())
                meters["v_loss"].update(loss.item())
                meters["v_acc_right"].update(right_acc.item())
                meters["v_acc_left"].update(left_acc.item())

        print(f"Epoch {i_epoch}: {get_summary(meters)}")

        if meters["v_loss"].avg < best_val_loss:
            best_val_loss = meters["v_loss"].avg
            print(f"Best Loss @ {i_epoch} (v_loss: {best_val_loss})")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument("--path", type=Path, default=Path("./dataset"))
    parser.add_argument("--lr", type=float, default=1e-3)
    parser.add_argument("--batch_size", type=int, default=256)
    parser.add_argument("--seed", type=int, default=10)
    parser.add_argument("--n_epochs", type=int, default=100)
    parser.add_argument("--pcgrad", action="store_true")
    parser.add_argument("--download", action="store_true")

    args = parser.parse_args()

    try:
        main(args)

    except Exception as e:
        raise e