<gh_stars>1-10
import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
from torchvision.utils import make_grid

import numpy as np
from tqdm import tqdm

from utils import AverageMeter, accuracy_top1

cudnn.benchmark = True


def _model_loop(args, loop_type, loader, model, opt, epoch, writer, adv, plot_gap=np.inf):
    losses = AverageMeter()
    acc_top1 = AverageMeter()

    is_train = (loop_type == 'Train')
    model = model.train() if is_train else model.eval()
    classifier_criterion = nn.CrossEntropyLoss()

    iterator = tqdm(enumerate(loader), total=len(loader))
    for i, (inp, target) in iterator:
        inp = inp.cuda(non_blocking=True)
        target = target.cuda(non_blocking=True)

        logits, inp_adv = model(inp, make_adv=adv, target=target)

        classifier_loss = classifier_criterion(logits, target)
        losses.update(classifier_loss.item(), inp.size(0))

        prec = accuracy_top1(logits, target)
        acc_top1.update(prec, inp.size(0))

        if is_train:
            opt.zero_grad()
            classifier_loss.backward()
            opt.step()
        if i % plot_gap == 0 and writer is not None:
            if adv:
                input_output = [inp[:8], inp_adv[:8], inp[:8] - inp_adv[:8],
                                inp[8:16], inp_adv[8:16], inp[8:16] - inp_adv[8:16]]
                input_output = torch.cat(input_output)
                input_output = make_grid(input_output, nrow=8, normalize=True, scale_each=True)
                writer.add_image(loop_type + '_adversarial', input_output, epoch + i // plot_gap)

        desc = ('{} Epoch:{} | Loss {:.4f} | Accuracy {:.3f} ||'
                .format(loop_type, epoch, losses.avg, acc_top1.avg))
        iterator.set_description(desc)
        iterator.refresh()

    if writer is not None:
        prec_type = 'adv' if adv else 'nat'
        descs = ['loss', 'accuracy']
        vals = [losses, acc_top1]
        for d, v in zip(descs, vals):
            writer.add_scalar('_'.join([prec_type, loop_type, d]), v.avg, epoch)

    return losses.avg, acc_top1.avg


def train_model(args, train_loader, val_loader, model, optimizer, schedule, writer):
    start_epoch, best_loss, best_acc = 0, np.inf, 0

    for epoch in range(start_epoch, args.max_epoch):
        train_loss, train_acc = _model_loop(args, 'Train', train_loader, model, optimizer, epoch, writer, args.adv_train)

        snapshot = {
            'model': model.state_dict(),
            'epoch': epoch,
        }

        last_epoch = (epoch == (args.max_epoch - 1))
        should_log = (epoch % args.log_gap == 0)
        if should_log or last_epoch:
            with torch.no_grad():
                test_loss, test_acc = _model_loop(args, 'Val', val_loader, model, None, epoch, writer, adv=False)
            if args.adv_train:
                adv_loss, adv_acc = _model_loop(args, 'Val', val_loader, model, None, epoch, writer, adv=True)

            our_acc = adv_acc if args.adv_train else test_acc
            is_best = our_acc > best_acc
            best_acc = max(our_acc, best_acc)

            torch.save(snapshot, args.model_save_path)
            if is_best:
                torch.save(snapshot, args.model_best_path)
        schedule.step()
    return model


def eval_model(args, test_loader, model, writer):
    test_loss, test_acc = _model_loop(args, 'Val', test_loader, model, None, 0, writer, adv=False, plot_gap=args.plot_gap)
    adv_loss, adv_acc = _model_loop(args, 'Val', test_loader, model, None, 0, writer, adv=True, plot_gap=args.plot_gap)

