<gh_stars>1-10
import dill
import tqdm

import torch
import torch.nn.functional as F
from torchvision import datasets, transforms

import apex.amp as amp


device = torch.device('cuda:0') if torch.cuda.device_count() > 0 else torch.device('cpu')

cifar10_mean = (0.4914, 0.4822, 0.4465)
cifar10_std = (0.2471, 0.2435, 0.2616)

mu = torch.tensor(cifar10_mean).view(3, 1, 1).to(device)
std = torch.tensor(cifar10_std).view(3, 1, 1).to(device)

upper_limit = ((1. - torch.zeros_like(mu)) / torch.ones_like(std))
lower_limit = ((0. - torch.zeros_like(mu)) / torch.ones_like(std))


def clamp(X, lower_limit, upper_limit):
    return torch.max(torch.min(X, upper_limit), lower_limit)


def get_loaders(dir_, batch_size, workers=4):
    train_transform = transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
    ])  # normalization part of model forward-pass
    test_transform = transforms.Compose([
        transforms.ToTensor(),
    ])  # normalization part of model forward-pass

    train_dataset = datasets.CIFAR10(
        dir_, train=True, transform=train_transform, download=True)
    test_dataset = datasets.CIFAR10(
        dir_, train=False, transform=test_transform, download=True)
    train_loader = torch.utils.data.DataLoader(
        dataset=train_dataset,
        batch_size=batch_size,
        shuffle=True,
        pin_memory=True,
        num_workers=workers,
    )
    test_loader = torch.utils.data.DataLoader(
        dataset=test_dataset,
        batch_size=batch_size,
        shuffle=False,
        pin_memory=True,
        num_workers=workers,
    )
    return train_loader, test_loader


def attack_pgd(model, X, y, epsilon, alpha, attack_iters, restarts, device, opt=None):
    max_loss = torch.zeros(y.shape[0]).to(device)
    max_delta = torch.zeros_like(X).to(device)
    for zz in range(restarts):
        delta = torch.zeros_like(X).to(device)
        for i in range(len(epsilon)):
            delta[:, i, :, :].uniform_(-epsilon[i][0][0].item(), epsilon[i][0][0].item())
        delta.data = clamp(delta, lower_limit - X, upper_limit - X)
        delta.requires_grad = True
        for _ in range(attack_iters):
            output = model(X + delta)
            index = torch.where(output.max(1)[1] == y)
            if len(index[0]) == 0:
                break
            loss = F.cross_entropy(output, y)
            if opt is not None:
                with amp.scale_loss(loss, opt) as scaled_loss:
                    scaled_loss.backward()
            else:
                loss.backward()
            grad = delta.grad.detach()
            d = delta[index[0], :, :, :]
            g = grad[index[0], :, :, :]
            d = clamp(d + alpha * torch.sign(g), -epsilon, epsilon)
            d = clamp(d, lower_limit - X[index[0], :, :, :], upper_limit - X[index[0], :, :, :])
            delta.data[index[0], :, :, :] = d
            delta.grad.zero_()
        all_loss = F.cross_entropy(model(X + delta), y, reduction='none').detach()
        max_delta[all_loss >= max_loss] = delta.detach()[all_loss >= max_loss]
        max_loss = torch.max(max_loss, all_loss)
    return max_delta


def evaluate_pgd(test_loader, model, attack_iters, restarts, device):
    epsilon = (8 / 255.) / torch.ones_like(std)
    alpha = (2 / 255.) / torch.ones_like(std)
    pgd_loss = 0
    pgd_acc = 0
    n = 0
    model.eval()
    pbar = tqdm.tqdm(total=len(test_loader), leave=False)
    pbar.set_description('PGD Evaluation')
    for i, (X, y) in enumerate(test_loader):
        X, y = X.to(device), y.to(device)
        pgd_delta = attack_pgd(model, X, y, epsilon, alpha, attack_iters, restarts, device)
        with torch.no_grad():
            output = model(X + pgd_delta)
            loss = F.cross_entropy(output, y)
            pgd_loss += loss.item() * y.size(0)
            pgd_acc += (output.max(1)[1] == y).sum().item()
            n += y.size(0)
        pbar.update(1)
    pbar.close()

    return pgd_loss / n, pgd_acc / n


def evaluate_standard(test_loader, model, device):
    test_loss = 0
    test_acc = 0
    n = 0
    model.eval()
    pbar = tqdm.tqdm(total=len(test_loader))
    pbar.set_description('Standard Evaluation')
    with torch.no_grad():
        for i, (X, y) in enumerate(test_loader):
            X, y = X.to(device), y.to(device)
            output = model(X)
            loss = F.cross_entropy(output, y)
            test_loss += loss.item() * y.size(0)
            test_acc += (output.max(1)[1] == y).sum().item()
            n += y.size(0)
            pbar.update(1)
    pbar.close()

    return test_loss / n, test_acc / n


class InputNormalize(torch.nn.Module):
    '''
    A module (custom layer) for normalizing the input to have a fixed 
    mean and standard deviation (user-specified).
    '''
    def __init__(self, new_mean, new_std):
        super(InputNormalize, self).__init__()
        new_std = new_std[..., None, None]
        new_mean = new_mean[..., None, None]

        self.register_buffer("new_mean", new_mean)
        self.register_buffer("new_std", new_std)

    def forward(self, x):
        x = torch.clamp(x, 0, 1)
        x_normalized = (x - self.new_mean)/self.new_std
        return x_normalized


class ModelwithInputNormalization(torch.nn.Module):
    def __init__(self, net, mean, std):
        super(ModelwithInputNormalization, self).__init__()
        self.normalizer = InputNormalize(mean, std)
        self.net = net

    def forward(self,  inp, with_latent=False, fake_relu=False, no_relu=False):
        normalized_inp = self.normalizer(inp)

        if no_relu and (not with_latent):
            print("WARNING: 'no_relu' has no visible effect if 'with_latent is False.")
        if no_relu and fake_relu:
            raise ValueError("Options 'no_relu' and 'fake_relu' are exclusive")

        output = self.net(normalized_inp, with_latent=with_latent,
                          fake_relu=fake_relu, no_relu=no_relu)

        return output


def load_checkpoint(load_path, mode, device):
    assert mode in ['train', 'eval']

    checkpoint = torch.load(load_path, pickle_module=dill, map_location=device)
    state_dict_path = "model"
    if not ("model" in checkpoint):
        state_dict_path = "state_dict"

    sd = checkpoint[state_dict_path]
    sd = {k.replace("module.model.", ""):v for k,v in sd.items()\
          if ("attacker" not in k) and ("normalizer" not in k)}
    # NOTE: fix for loading older checkpoints
    sd = {k.replace("net.", ""): v for k, v in sd.items()}

    load_epoch = checkpoint['epoch']
    optim_sd = None
    best_acc = None
    r_time_avg = None
    if mode == 'train':
        optim_sd = checkpoint["optimizer_state_dict"]
        best_acc = checkpoint["best_acc"]
        r_time_avg = checkpoint["time_per_epoch"]

    # clear out memory
    del checkpoint

    ckpt = {'model_sd': sd,
            'load_epoch': load_epoch,
            'optim_sd': optim_sd,
            'best_acc': best_acc,
            'r_time_avg': r_time_avg}

    return ckpt
