train_utils.py

import os, sys, copy
from os.path import dirname, realpath
import torch
import tqdm.notebook as tq
from efficientnet_pytorch import EfficientNet
import torch.nn as nn
import torch.optim as optim
import numpy as np

# LS >>> what is this, winston?
sys.path.append(dirname(realpath(__file__))) 

from sam.sam import SAM
from eval_utils import simple_accuracy

def get_SGD(model_params, lr, momentum, decay):
    return optim.SGD(model_params, lr=lr, momentum=momentum,
                     weight_decay=decay)


def get_Adam(model_params, lr, momentum, decay):
    return optim.Adam(model_params, lr=lr, weight_decay=decay)


def get_SAM(model_params, lr, momentum, decay):
    base_optimizer = torch.optim.SGD
    return SAM(model_params, base_optimizer, lr=lr, momentum=momentum,
               weight_decay=decay)

def get_lossWeights(beta, num_classes, data_dict):
    # >>> begin getting weights

    # get dictionary of training data [now expected as parameter]

    # convert to dictionary of quantity of training data per class
    label_to_quant = {key//2:len(data_dict[key]['annotations']) for key in data_dict}

    # compute weights based on
    # https://arxiv.org/pdf/1901.05555.pdf
    # https://towardsdatascience.com/handling-class-imbalanced-data-using-a-loss-specifically-made-for-it-6e58fd65ffab
    effective_num = np.zeros(num_classes)
    for i in label_to_quant:
        effective_num[i] = label_to_quant[i]

    effective_num = 1 - np.power(beta, effective_num)
    weights = (1-beta)/effective_num
    weights = weights/np.sum(weights) * num_classes
    weights = torch.FloatTensor(weights)
    return weights


def get_class_type_loss_weight(beta, num_class_types, effective_num):
    effective_num = 1 - np.power(beta, effective_num)
    weights = (1-beta)/effective_num
    weights = weights/np.sum(weights)*num_class_types
    weights = torch.FloatTensor(weights)
    return weights


def get_class_type_count(data_n_arr, class_type_arr, num_types):
    class_type_n_arr = [0 for i in range(num_types)]
    for idx, class_type in enumerate(class_type_arr):
        class_type_n_arr[class_type] += data_n_arr[idx]
    return class_type_n_arr


def get_model(device, num_classes, net_str, optim_type, model=None, lr=0.001,
              momentum=0.9, decay=0.0005, loss_weight=None):
    """optim_type == SGD | Adam | SAM"""
    hidden_count_dict = {
        "efficientnet-b0": 1280,
        "efficientnet-b1": 1280,
        "efficientnet-b2": 1408,
        "efficientnet-b3": 1536,
        "efficientnet-b4": 1792,
        "efficientnet-b5": 2048,
        "efficientnet-b6": 2304,
        "efficientnet-b7": 2560,
        "efficientnet-b8": 2816,
        "efficientnet-l2": 5504
    }
    if model is None:
        model = EfficientNet.from_name(net_str)
        model._fc = nn.Linear(hidden_count_dict[net_str], num_classes)
    criterion = nn.CrossEntropyLoss(weight=loss_weight)
    optim_dict = {"SGD": get_SGD, "Adam": get_Adam, "SAM": get_SAM}
    assert optim_type in optim_dict, "invalid optim_type"
    optimizer = optim_dict[optim_type](model.parameters(), lr, momentum, decay)
    model.to(device)
    criterion.to(device)

    return model, criterion, optimizer


def train(checkpoint_dir, net, train_loader, vali_loader, net_str, data_label,
          rho, device, criterion, optimizer, beta=0.0, epochs=1, div=True):
    losses = []
    vali_losses = []
    vali_accues = []

    # write save dir
    if net_str == "efficientnet-b0":
        net_str = "EfficientNet"
    loss_str = type(criterion).__name__
    optim_str = type(optimizer).__name__
    lr = optimizer.param_groups[0]['lr']
    momentum = optimizer.param_groups[0]["momentum"]
    gamma = optimizer.param_groups[0]['weight_decay']
    assert lr is not None and gamma is not None, "optimizer parameter error"
    print(f"training params: net={net_str}\tloss={loss_str}\toptim={optim_str}"
          f"\tepochs={epochs}\tlr={lr}\tweight_decay={float(gamma)}")

    def url_func(epo):
        if momentum == 0.9:
            return os.path.join(
                checkpoint_dir,
                f"{net_str}_{loss_str}_beta{beta}_{optim_str}_lr{lr}_gamma{gamma}_"
                f"{data_label}_rho{rho}_epoch{epo}.pt")
        else:
            return os.path.join(
                checkpoint_dir,
                f"{net_str}_{loss_str}_beta{beta}_{optim_str}_lr{lr}_gamma{gamma}_"
                f"momentum{momentum}_{data_label}_rho{rho}_epoch{epo}.pt")

    # load in saved checkpoints
    start_epoch = 0
    for i in range(epochs):
        curr_epoch = epochs-i
        curr_url = url_func(curr_epoch)
        if os.path.exists(curr_url):
            print(f"load checkpoint from={curr_url}")
            state = torch.load(curr_url)
            net.load_state_dict(state['net'])
            optimizer.load_state_dict(state['optimizer'])
            start_epoch = state['epoch']
            vali_losses = state['vali_losses']
            vali_accues = state['vali_accues']
            break

    epoch = start_epoch
    print("File Path:", url_func(start_epoch))
    for epoch in range(start_epoch, epochs):
        net.train()
        t = tq.tqdm(enumerate(train_loader), desc=f"train epoch={epoch}",
                    position=0, leave=True, total=len(train_loader))
        for i, batch in t:
            inputs, labels = batch
            # ***map to network output, specific for CIFAR-50***
            if div:
                labels = labels//2

            def closure():
                loss = criterion(net(inputs.to(device)),
                                 (labels).to(device))
                loss.backward()
                return loss
            # zero the parameter gradients
            optimizer.zero_grad()
            # forward + backward + optimize
            loss = criterion(net(inputs.to(device)), (labels).to(device))
            loss.backward()
            if optim_str == "SAM":
                optimizer.step(closure)
            else:
                optimizer.step()
            losses.append(loss.item())
            t.set_description(f"loss={loss.item():.3f}\tepoch={epoch}")

        # validation
        net.eval()
        vali_accuracy = simple_accuracy(net, vali_loader, device, div=div)
        print("epoch {}\tloss={:.3f}\taccuracy={:.3f}".format(
            epoch, losses[-1], vali_accuracy))
        vali_losses.append(losses[-1])
        vali_accues.append(vali_accuracy)

        # save state_dict
        state = {'epoch': epoch+1, 'net': net.state_dict(),
                 'optimizer': optimizer.state_dict(),
                 'vali_losses': vali_losses,
                 'vali_accues': vali_accues}
        curr_save_url = url_func(epoch+1)
        prev_save_url = url_func(epoch)
        print(f"save model checkpoint at={curr_save_url}")
        if not os.path.exists(checkpoint_dir):
            os.makedirs(checkpoint_dir)
        torch.save(state, curr_save_url)
        if os.path.exists(prev_save_url) and epoch % 10 != 0:
            os.remove(prev_save_url)

    # net.load_state_dict(best_state['net'])
    return net, vali_losses, vali_accues


def load_cRT_model(root_dir, device, net_str, loss_str, optim_str, rho, lr,
                   gamma, beta, epochs, num_classes, optim_type, data_label, 
                   cRT=True, resampled=True):
    checkpoint_dir = os.path.join(root_dir, "checkpoints")
    hidden_count_dict = {
        "efficientnet-b0": 1280,
        "efficientnet-b1": 1280,
        "efficientnet-b2": 1408,
        "efficientnet-b3": 1536,
        "efficientnet-b4": 1792,
        "efficientnet-b5": 2048,
        "efficientnet-b6": 2304,
        "efficientnet-b7": 2560,
        "efficientnet-b8": 2816,
        "efficientnet-l2": 5504
    }
    if net_str == "efficientnet-b0":
        cp_net_str = "EfficientNet"
    else:
        cp_net_str = net_str
    model_base = f"{cp_net_str}_{loss_str}_beta{beta}_{optim_str}_lr{lr}_gamma{gamma}_" \
                 f"{data_label}_rho{rho}_epoch{epochs}"
    model_url = os.path.join(checkpoint_dir, f"{model_base}.pt")
    if not os.path.exists(model_url):
        print(f"model url: {model_url} doesn't exist")
        return
    print(f"load model from={model_url}")
    print(f"model params: net={net_str}\tloss={loss_str}\toptim={optim_str}"
          f"\tepochs={epochs}\tlr={lr}\tweight_decay={float(gamma)}")
    state = torch.load(model_url)
    model = EfficientNet.from_name(net_str)
    model._fc = nn.Linear(hidden_count_dict[net_str], num_classes)
    model.load_state_dict(state['net'])

    if resampled:
        cRT_folder = os.path.join(checkpoint_dir, f"{model_base}_cRT")
    else:
        cRT_folder = os.path.join(
            checkpoint_dir, f"{model_base}_resampled0_cRT")
    if cRT:
        if not os.path.exists(cRT_folder):
            os.makedirs(cRT_folder)

        # fix representation & randomize classifier
        for param in model.parameters():
            param.requires_grad = False
        model._fc = nn.Linear(hidden_count_dict[net_str], num_classes) \
            .to(device)
        model._dropout = nn.Dropout(p=0.2, inplace=False).to(device)
        model._avg_pooling = nn.AdaptiveAvgPool2d(output_size=1).to(device)
        model._bn1 = nn.BatchNorm2d(
            hidden_count_dict[net_str], eps=0.001,
            momentum=0.010000000000000009,
            affine=True, track_running_stats=True).to(device)
        model.train()
    return model, cRT_folder


def tau_norm(net, tau, num_classes, device):
    tau_net = copy.deepcopy(net)
    weights = list(tau_net._fc.parameters())[0].data.clone()
    normB = torch.norm(weights, 2, 1)
    ws = weights.clone()
    for i in range(weights.size(0)):
        ws[i] = ws[i] / torch.pow(normB[i], tau)
    list(tau_net._fc.parameters())[0].data = ws.to(device)
    list(tau_net._fc.parameters())[1].data = torch.zeros(num_classes).to(device)

    return tau_net