mart.py

from __future__ import print_function
import argparse
import os
import time
import warnings
import models
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
from torch.autograd import Variable
from torch.utils.tensorboard import SummaryWriter
from torchvision import datasets, transforms
from utils import Logger, random_seed, save_checkpoint

warnings.filterwarnings("ignore")

parser = argparse.ArgumentParser(description="Adversarial defense with projection removal")
parser.add_argument(
    "--arch", type=str, default="ResNet18", choices=["WideResNet", "ResNet18"]
)
# data
parser.add_argument(
    "--data", type=str, default="CIFAR10", choices=["CIFAR10", "CIFAR100", "SVHN"]
)
parser.add_argument(
    "--data-path", type=str, default="~/datasets/", help="where is the dataset CIFAR-10"
)
parser.add_argument(
    "--batch-size",
    type=int,
    default=128,
    metavar="N",
    help="input batch size for training (default: 128)",
)
parser.add_argument(
    "--test-batch-size",
    type=int,
    default=1000,
    metavar="N",
    help="input batch size for testing (default: 128)",
)
# traning setting
parser.add_argument(
    "--epochs", type=int, default=120, metavar="N", help="number of epochs to train"
)
parser.add_argument("--weight-decay", "--wd", default=3.5e-3, type=float, metavar="W")
parser.add_argument(
    "--lr", type=float, default=0.01, metavar="LR", help="learning rate"
)
parser.add_argument(
    "--momentum", type=float, default=0.9, metavar="M", help="SGD momentum"
)
parser.add_argument(
    "--no-cuda", action="store_true", default=False, help="disables CUDA training"
)
parser.add_argument(
    "--norm",
    default="l_inf",
    type=str,
    choices=["l_inf", "l_2"],
    help="The threat model",
)
parser.add_argument("--epsilon", default=8.0 / 255, type=eval, help="perturbation")
parser.add_argument("--num-steps", default=10, type=int, help="perturb number of steps")
parser.add_argument(
    "--step-size", default=2.0 / 255, type=eval, help="perturb step size"
)
parser.add_argument(
    "--beta", default=5.0, type=float, help="regularization, i.e., 1/lambda in TRADES"
)
# Eval PGD setting
parser.add_argument("--test-epsilon", default=8.0 / 255, type=eval, help="perturbation")
parser.add_argument(
    "--test-step-size", default=2/ 255, type=eval, help="perturb step size"
)
parser.add_argument(
    "--test-num-steps", default=20, type=int, help="perturb number of steps"
)
# resume
parser.add_argument(
    "--start-epoch", type=int, default=1, metavar="N", help="retrain from which epoch"
)
parser.add_argument(
    "--resume_path", default="", type=str, help="directory of model for retraining"
)
# save checkpoint
parser.add_argument(
    "--result-dir",
    default="results/MART",
    help="directory of model for saving checkpoint",
)
parser.add_argument(
    "--save-freq", "-s", default=1, type=int, metavar="N", help="save frequency"
)
parser.add_argument(
    "--seed", type=int, default=1, metavar="S", help="random seed (default: 1)"
)
parser.add_argument(
    "--lmbda", type=float, default=0.001, help="lambda coeff for removing projection"
)


args = parser.parse_args()

if args.data == "CIFAR100":
    NUM_CLASSES = 100
elif args.data=="CIFAR10" or args.data=="SVHN":
    NUM_CLASSES = 10


if args.seed is not None:
    random_seed(args.seed)


args_path = (
    "epoch"
    + str(args.epochs)
    + "_bs"
    + str(args.batch_size)
    + "_lr"
    + str(args.lr)
    + "_wd"
    + str(args.weight_decay)
    + "_eps"
    + str(args.epsilon)
    + "_norm"
    + str(args.norm)
    + "_beta"
    + str(args.beta)
    + "_lmbda"
    + str(args.lmbda)
    + "test-step-size"
    + str(args.test_step_size)
)

checkpoint_path = os.path.join(
    args.result_dir, args.data, args.arch, args_path, "checkpoints"
)

if not os.path.exists(checkpoint_path):
    os.makedirs(checkpoint_path)

writer = SummaryWriter(
    os.path.join(args.result_dir, args.data, args.arch, args_path, "tensorboard_logs")
)

logger = Logger(
    os.path.join(args.result_dir, args.data, args.arch, args_path, "output.log")
)

best_nature_acc = 0
best_robust_acc = 0

use_cuda = not args.no_cuda and torch.cuda.is_available()
device = torch.device("cuda:0" if use_cuda else "cpu")
kwargs = {"num_workers": 1, "pin_memory": True} if use_cuda else {}
print(f"Using device: {device} (CUDA device index: {torch.cuda.current_device()})")
print(f"Device Name: {torch.cuda.get_device_name(torch.cuda.current_device())}")
print(f"num classes = {NUM_CLASSES}")


# setup data loader
transform_train = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    ])
transform_test = transforms.Compose(
    [
        transforms.ToTensor(),
    ]
)


if args.data == "SVHN":
    trainset = torchvision.datasets.SVHN(root=args.data_path, split="train", download=False, transform=transform_test)
    train_loader = torch.utils.data.DataLoader(trainset, batch_size=args.batch_size, shuffle=False, **kwargs)
    eval_trainset = torchvision.datasets.SVHN(args.data_path, split="train", download=False, transform=transform_test)
    eval_train_loader = torch.utils.data.DataLoader(eval_trainset, batch_size=args.test_batch_size, shuffle=False, **kwargs)
    testset = torchvision.datasets.SVHN(args.data_path, split="test", download=False, transform=transform_test)
    test_loader = torch.utils.data.DataLoader(testset, batch_size=args.test_batch_size, shuffle=False, **kwargs)

elif args.data=="CIFAR10" or args.data=="CIFAR100":
    print(getattr(datasets, args.data))
    trainset = getattr(datasets, args.data)(
      root=args.data_path, train=True, download=True, transform=transform_train
    )

    testset = getattr(datasets, args.data)(
      root=args.data_path, train=False, download=True, transform=transform_test
    )
    train_loader = torch.utils.data.DataLoader(
      trainset, batch_size=args.batch_size, shuffle=True, **kwargs
    )
    test_loader = torch.utils.data.DataLoader(
      testset, batch_size=args.test_batch_size, shuffle=False, **kwargs
    )


best_robust_path = checkpoint_path+"-best-robust.pth"
best_natural_path = checkpoint_path+"-best-natural.pth"

def save_checkpoint(state, epoch, is_best, model_type, save_path, save_freq):
    """
    Saves the checkpoint based on the provided state and conditions.
    
    Args:
    - state (dict): Contains model state, optimizer state, and other metrics.
    - epoch (int): Current epoch number.
    - is_best (bool): Whether this is the best model so far (based on accuracy).
    - model_type (str): Type of model, e.g., 'nature' or 'robust'.
    - save_path (str): Path to save the checkpoint.
    - save_freq (int): Frequency (in epochs) for saving checkpoints.
    """
    # Ensure directory exists
    os.makedirs(os.path.dirname(save_path), exist_ok=True)
    
    # Save checkpoint every save_freq epochs
    save_path+=f"{epoch}.pth"
    if epoch % save_freq == 0:
        torch.save(state, save_path)
        print(f"Checkpoint saved at {save_path}")

    global best_robust_path
    global best_natural_path
    # Save the best model checkpoint
    if is_best:
        #best_path = save_path.replace(".pth", f"_{model_type}_best.pth")
        if(model_type=='robust'): best_path = best_robust_path
        else: best_path = best_natural_path
        torch.save(state, best_path)
        print(f"Best model checkpoint saved at {best_path}")

def load_checkpoint(checkpoint_path, model, optimizer=None):
    """
    Loads a checkpoint and restores model and optimizer states.

    Args:
    - checkpoint_path (str): Path to the checkpoint file.
    - model (torch.nn.Module): Model to load the parameters into.
    - optimizer (torch.optim.Optimizer, optional): Optimizer to restore its state. Default is None.

    Returns:
    - checkpoint (dict): The loaded checkpoint.
    """
    if not os.path.exists(checkpoint_path):
        raise FileNotFoundError(f"Checkpoint not found at {checkpoint_path}")
    
    checkpoint = torch.load(checkpoint_path)
    model.load_state_dict(checkpoint["model_state_dict"])
    
    if optimizer is not None:
        optimizer.load_state_dict(checkpoint["opt_state_dict"])
    
    print(f"Checkpoint loaded from {checkpoint_path}")
    return checkpoint



def match_softmax_or_label(softmax_probs, labels):
    """
    Compares softmax predictions with labels and returns:
    - The softmax probabilities if the highest softmax index matches the label.
    - The one-hot encoded label if the prediction is incorrect.
    
    Args:
        softmax_probs (torch.Tensor): Softmax probabilities of shape (batch_size, num_classes).
        labels (torch.Tensor): Ground truth labels of shape (batch_size,).

    Returns:
        torch.Tensor: Modified probabilities based on correctness.
    """
    batch_size, num_classes = softmax_probs.shape

    # Get predicted class from softmax
    predicted_classes = torch.argmax(softmax_probs, dim=1)

    # Create one-hot encoded labels
    one_hot_labels = torch.zeros_like(softmax_probs)
    one_hot_labels.scatter_(1, labels.unsqueeze(1), 1.0)

    # Condition: Use softmax_probs when prediction is correct, else use one-hot labels
    mask = (predicted_classes == labels).unsqueeze(1).expand_as(softmax_probs)
    result = torch.where(mask, softmax_probs, one_hot_labels)

    return result

def cross_entropy_between_logits(logits1, logits2):
    
    log_p2 = torch.log_softmax(logits2, dim=1)  # Log probabilities from second case
    loss = -torch.sum(logits1 * log_p2, dim=1).mean()
    
    return loss


def mart_loss(
    model,
    model_teacher, epoch,
    x_natural,
    y,
    optimizer,
    step_size=0.007,
    epsilon=0.031,
    perturb_steps=10,
    beta=5.0,
    distance="l_inf",
):
    kl = nn.KLDivLoss(reduction="none")
    model.eval()
    batch_size = len(x_natural)
    # generate adversarial example
    x_adv = (
        x_natural.detach() + 0.001 * torch.randn(x_natural.shape).to(device).detach()
    )

    if distance == "l_inf":
        for _ in range(perturb_steps):
            x_adv.requires_grad_()
            with torch.enable_grad():
                loss_ce = F.cross_entropy(model(x_adv), y)

            grad = torch.autograd.grad(loss_ce, [x_adv])[0]
            x_adv = x_adv.detach() + step_size * torch.sign(grad.detach())
            x_adv = torch.min(
                torch.max(x_adv, x_natural - epsilon), x_natural + epsilon
            )
            x_adv = torch.clamp(x_adv, 0.0, 1.0)
    else:
        x_adv = torch.clamp(x_adv, 0.0, 1.0)
    model.train()

    x_adv = Variable(torch.clamp(x_adv, 0.0, 1.0), requires_grad=False)
    # zero gradient
    optimizer.zero_grad()

    logits = model(x_natural) 

    logits_adv = model(x_adv) 

    
    lmbda = args.lmbda
    closest_samples, closest_labels, closest_samples_second, closest_labels_second=find_closest_interclass_samples_with_labels(logits,y)
    logits = remove_projection(logits, closest_samples, closest_samples_second, lambda_coeff=lmbda)
    closest_samples, closest_labels, closest_samples_second, closest_labels_second = find_closest_interclass_samples_with_labels(logits_adv, y)
    logits_adv = remove_projection(logits_adv, closest_samples, closest_samples_second, lambda_coeff=lmbda)

    adv_probs = F.softmax(logits_adv, dim=1)

    tmp1 = torch.argsort(adv_probs, dim=1)[:, -2:]

    new_y = torch.where(tmp1[:, -1] == y, tmp1[:, -2], tmp1[:, -1])

    if epoch <=-1:
        logits_teacher = model_teacher((x_adv))
        y_ = match_softmax_or_label(F.softmax(logits_teacher,dim=1),y)
        adv_loss = cross_entropy_between_logits(y_, logits_adv)

        nat_loss=F.cross_entropy(logits, y)
        loss_adv = .25*(nat_loss) + adv_loss+  F.nll_loss(
            torch.log(1.0001 - adv_probs + 1e-12), new_y
        )
    else:
        loss_adv = F.cross_entropy(logits_adv, y) + F.nll_loss(
            torch.log(1.0001 - adv_probs + 1e-12), new_y
        )

    nat_probs = F.softmax(logits, dim=1)
    
    true_probs = torch.gather(nat_probs, 1, (y.unsqueeze(1)).long()).squeeze()

    loss_robust = (1.0 / batch_size) * torch.sum(
        torch.sum(kl(torch.log(adv_probs + 1e-12), nat_probs), dim=1)
        * (1.0000001 - true_probs)
    )

    loss = loss_adv + float(beta) * loss_robust

    return loss, logits, logits_adv

def find_closest_interclass_samples_with_labels(features, labels):
    """
    Find the closest inter-class sample and its class label for each sample in the batch.

    Args:
        features: Tensor of shape [batch_size, feature_dim], feature embeddings.
        labels: Tensor of shape [batch_size], class labels.

    Returns:
        closest_samples: Tensor of shape [batch_size, feature_dim], closest inter-class samples.
        closest_labels: Tensor of shape [batch_size], class labels of the closest inter-class samples.
    """
    batch_size = features.size(0)
    # print(torch.sum(features))

    # Check for NaN or Inf in features
    if torch.isnan(features).any() or torch.isinf(features).any():
        raise ValueError("Features tensor contains NaN or Inf values.")

    # Compute pairwise distances (L2 norm)
    pairwise_distances = torch.cdist(features, features, p=2)  # Shape: [batch_size, batch_size]

    # Create a mask for inter-class samples
    label_equal = labels.unsqueeze(1) == labels.unsqueeze(0)  # Shape: [batch_size, batch_size]
    interclass_mask = ~label_equal  # Invert mask: True for inter-class, False for intra-class

    # Set intra-class distances to Inf
    pairwise_distances = pairwise_distances.masked_fill(~interclass_mask, float('inf'))

    # Check for rows with no valid inter-class samples
    if torch.isinf(pairwise_distances).all(dim=1).any():
        raise ValueError("No valid inter-class samples for some inputs.")

    # Find the closest inter-class sample for each sample
    # closest_indices = torch.argmin(pairwise_distances, dim=1)  # Shape: [batch_size]
    min_indices = torch.topk(pairwise_distances, k=2, largest=False, dim=1).indices

    # Extract the min and second min indices separately
    min_index = min_indices[:, 0]
    second_min_index = min_indices[:, 1]
    closest_samples = features[min_index]  # Shape: [batch_size, feature_dim]
    closest_samples_second = features[second_min_index]  # Shape: [batch_size, feature_dim]
    closest_labels = labels[min_index]  # Shape: [batch_size]
    closest_labels_second = labels[second_min_index]

    return closest_samples, closest_labels, closest_samples_second, closest_labels_second

def remove_projection(x, y, y_second, lambda_coeff=None):
    """
    Removes the projection of y onto x to make x and y more dissimilar.

    Args:
        x (torch.Tensor): Anchor vector (batch_size, feature_dim)
        y (torch.Tensor): Distractor vector (batch_size, feature_dim)
        lambda_coeff (float): Scaling factor for projection removal

    Returns:
        torch.Tensor: Modified x with less similarity to y
    """
    # Compute projection of y onto x
    proj_coeff = (torch.sum(x * y, dim=1, keepdim=True) / (torch.sum(x * x, dim=1, keepdim=True)))
    proj_x_y = proj_coeff * x  # Projection vector

    # Subtract projection from x
    x_modified = x - lambda_coeff * proj_x_y
    return x_modified

def loss_cross_entropy_with_closest_class(features, labels, num_classes):
    """
    Compute the cross-entropy loss between features and the class of the closest inter-class sample.

    Args:
        features: Tensor of shape [batch_size, feature_dim], original feature embeddings.
        labels: Tensor of shape [batch_size], class labels.
        adversarial_samples: Tensor of shape [batch_size, feature_dim], adversarial samples.
        num_classes: Integer, the number of classes in the dataset.

    Returns:
        Loss value (scalar).
    """
    # Find closest inter-class samples and their labels
    closest_samples, closest_labels, closest_samples_second, closest_labels_second = find_closest_interclass_samples_with_labels(features, labels)


    # Normalize logits and compute cross-entropy loss
    lmbda = args.lmbda
    features = remove_projection(features, closest_samples, lambda_coeff=lmbda)
    loss = F.cross_entropy(features, labels)
    # loss = cross_entropy_between_logits(closest_samples,features)

    # Maximize the cross-entropy loss (minimize the negative of it)
    return loss

def train(args, model,model_teacher, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)

        optimizer.zero_grad()

        # calculate robust loss
        loss, logits, logits_adv= mart_loss(
            model=model,
            model_teacher=model_teacher,epoch=epoch,
            x_natural=data,
            y=target,
            optimizer=optimizer,
            step_size=args.step_size,
            epsilon=args.epsilon,
            perturb_steps=args.num_steps,
            beta=args.beta,
            distance=args.norm,
        )

        loss.backward()  # Retain graph if needed
        optimizer.step()


def adjust_learning_rate(optimizer, epoch, mode):
    """decrease the learning rate"""
    lr = args.lr

    if mode=='120e':
      if epoch >= 100:
          lr = args.lr * 0.001
      elif epoch >= 90:
          lr = args.lr * 0.01
      elif epoch >= 75:
          lr = args.lr * 0.1

    elif mode=='80e':
      if epoch >= 50:
        lr = args.lr * 0.1
      if epoch >= 65:
        lr = args.lr * 0.01

    for param_group in optimizer.param_groups:
        param_group["lr"] = lr


def _pgd_whitebox(
    model, X,y, epsilon=0.031, step_size=0.003, num_steps=20, random=True
):
    out = model(X)
    err = (out.data.max(1)[1] != y.data).float().sum()
    X_pgd = Variable(X.data, requires_grad=True)
    if random:
        random_noise = (
            torch.FloatTensor(*X_pgd.shape).uniform_(-epsilon, epsilon).to(device)
        )
        X_pgd = Variable(X_pgd.data + random_noise, requires_grad=True)

    for _ in range(num_steps):
        opt = optim.SGD([X_pgd], lr=1e-3)
        opt.zero_grad()

        with torch.enable_grad():
            loss = nn.CrossEntropyLoss()(model(X_pgd), y)
        loss.backward()
        eta = step_size * X_pgd.grad.data.sign()
        X_pgd = Variable(X_pgd.data + eta, requires_grad=True)
        eta = torch.clamp(X_pgd.data - X.data, -epsilon, epsilon)
        X_pgd = Variable(X.data + eta, requires_grad=True)
        X_pgd = Variable(torch.clamp(X_pgd, 0, 1.0), requires_grad=True)
    err_pgd = (model(X_pgd).data.max(1)[1] != y.data).float().sum()

    return err, err_pgd


def eval_adv_whitebox(
    model,model_teacher, device, test_loader, epsilon, step_size, num_steps, random
):
    """
    evaluate model by white-box attack
    """
    model.eval()
    robust_err_total = 0
    natural_err_total = 0
    count = 0
    for data, target in test_loader:
        data, target = data.to(device), target.to(device)
        # pgd attack
        X, y = Variable(data, requires_grad=True), Variable(target)
        count += X.shape[0]
        err_natural, err_robust = _pgd_whitebox(
            model, X, y, epsilon, step_size, num_steps, random
        )
        robust_err_total += err_robust
        natural_err_total += err_natural

    nature_acc = 1.0 - (natural_err_total / count)
    robust_acc = 1.0 - (robust_err_total / count)

    return nature_acc, robust_acc


def main():
    global best_nature_acc, best_robust_acc

    logger.info(args)

    # load model
    print(dir(models))

    model = getattr(models, args.arch)(num_classes=NUM_CLASSES)
    #print(model)

    # dataparallel
    model = model.to(device)
    
    model_teacher =  getattr(models, args.arch)(num_classes=NUM_CLASSES)
    model_teacher=model_teacher.to(device)
    model_teacher = model_teacher.eval()

    optimizer = optim.SGD(
        model.parameters(),
        lr=args.lr,
        momentum=args.momentum,
        weight_decay=args.weight_decay,
    )

    start_epoch = 1

    if args.start_epoch > 1:
        logger.info("Retrain from epoch %d", (args.start_epoch))
        state_dict = torch.load(args.resume_path, map_location=device)
        optimizer.load_state_dict(state_dict["opt_state_dict"])
        model.load_state_dict(state_dict["model_state_dict"])

    logger.info("Epoch \t Time \t Test ACC \t Test Robust ACC")

    mode = ''
    if args.epochs == 120: mode = '120e'
    elif args.epochs == 80: mode = '80e'

    for epoch in range(start_epoch, args.epochs + 1):
        # adjust learning rate for SGD
        adjust_learning_rate(optimizer, epoch, mode)

        # train
        start_time = time.time()

        train(args, model,model_teacher,device, train_loader, optimizer, epoch)

        # eval
        nature_acc, robust_acc = eval_adv_whitebox(
            model,
            model_teacher,
            device,
            test_loader,
            args.test_epsilon,
            args.test_step_size,
            args.test_num_steps,
            True,
        )
        epoch_time = time.time() - start_time

        logger.info(
            "%d\t %d \t %.4f \t %.4f", epoch, epoch_time, nature_acc, robust_acc
        )

        writer.add_scalar("Test/Acc", float(nature_acc), epoch)
        writer.add_scalar("Test/Robust Acc", float(robust_acc), epoch)

        # Save checkpoint
        is_best = nature_acc > best_nature_acc
        best_nature_acc = max(nature_acc, best_nature_acc)
        
        save_checkpoint(
            {
                "epoch": epoch,
                "model_state_dict": model.state_dict(),
                "opt_state_dict": optimizer.state_dict(),
                "nature_acc": float(nature_acc),
                "robust_acc": float(robust_acc),
            },
            epoch,
            is_best,
            "nature",
            save_path=checkpoint_path,
            save_freq=args.save_freq,
        )

        is_best_robust = robust_acc > best_robust_acc
        best_robust_acc = max(robust_acc, best_robust_acc)
        save_checkpoint(
            {
                "epoch": epoch,
                "model_state_dict": model.state_dict(),
                "opt_state_dict": optimizer.state_dict(),
                "nature_acc": float(nature_acc),
                "robust_acc": float(robust_acc),
            },
            epoch,
            is_best_robust,
            "robust",
            save_path=checkpoint_path,
            save_freq=args.save_freq,
        )

    logger.info("Best Nature ACC %.4f", best_nature_acc)
    logger.info("Best Robust ACC %.4f", best_robust_acc)
    writer.close()


if __name__ == "__main__":
    main()