cs291a_hw1.py

# -*- coding: utf-8 -*-
"""cs291A-hw1.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1JjGXaZmtAYQgC1Fn-AUzJUmtCHtKJqmw

# data_utils.py
"""

from torchvision.datasets import CIFAR10
from torch.utils.data import DataLoader, Subset
import torch
import numpy as np
import os
from torchvision import transforms

class NormalizeByChannelMeanStd(torch.nn.Module):
    def __init__(self, mean, std):
        super(NormalizeByChannelMeanStd, self).__init__()
        if not isinstance(mean, torch.Tensor):
            mean = torch.tensor(mean)
        if not isinstance(std, torch.Tensor):
            std = torch.tensor(std)
        self.register_buffer("mean", mean)
        self.register_buffer("std", std)

    def forward(self, tensor):
        return self.normalize_fn(tensor, self.mean, self.std)

    def extra_repr(self):
        return 'mean={}, std={}'.format(self.mean, self.std)

    def normalize_fn(self, tensor, mean, std):
        """Differentiable version of torchvision.functional.normalize"""
        # here we assume the color channel is in at dim=1
        mean = mean[None, :, None, None]
        std = std[None, :, None, None]
        return tensor.sub(mean).div(std)

def cifar10_dataloader(batch_size=64, data_dir='./data/', val_ratio=0.1):
    train_transform = transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
    ])

    test_transform = transforms.Compose([
        transforms.ToTensor(),
    ])

    train_size = int(50000 * (1 - val_ratio))
    val_size = 50000 - train_size

    train_set = Subset(CIFAR10(data_dir, train=True, transform=train_transform, download=True), list(range(train_size)))
    val_set = Subset(CIFAR10(data_dir, train=True, transform=test_transform, download=True),
                     list(range(train_size, train_size + val_size)))
    test_set = CIFAR10(data_dir, train=False, transform=test_transform, download=True)

    train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=2, pin_memory=True)
    val_loader = DataLoader(val_set, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True)
    test_loader = DataLoader(test_set, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True)

    dataset_normalization = NormalizeByChannelMeanStd(
        mean=[0.4914, 0.4822, 0.4465], std=[0.2470, 0.2435, 0.2616])
    return train_loader, val_loader, test_loader, dataset_normalization

"""# model_util.py"""

import time
import torch
import torch.nn as nn

class BasicModule(nn.Module):
    def __init__(self):
        super(BasicModule, self).__init__()
        self.model_name = str(type(self))

    def load(self, path, map_location=None):
        self.load_state_dict(torch.load(path, map_location))

    def save(self, name=None):
        if name is None:
            prefix = 'checkpoints/' + self.model_name + '_'
            name = time.strftime(prefix + '%m%d_%H:%M:%S.pth')
        torch.save(self.state_dict(), name)
        return name

    def no_grad(self):
        for param in self.parameters():
            param.requires_grad = False

    def with_grad(self):
        for param in self.parameters():
            param.requires_grad = True

    def clear_grad(self):
        for param in self.parameters():
            param.grad = None

class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, in_planes, planes, stride=1, activation_fn=nn.ReLU()):
        super(BasicBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)
        self.activation_fn = activation_fn
        self.shortcut = nn.Sequential()
        if stride != 1 or in_planes != self.expansion * planes:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(self.expansion * planes)
            )

    def forward(self, x):
        out = self.activation_fn(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out += self.shortcut(x)
        out = self.activation_fn(out)
        return out

class ResNet(BasicModule):
    def __init__(self, block, num_blocks, num_classes=10, activation_fn=nn.ReLU, conv1_size=3):
        super(ResNet, self).__init__()
        self.in_planes = 64
        self.activation_fn = activation_fn(beta=10) if activation_fn == nn.Softplus else activation_fn()

        kernel_size, stride, padding = {3: [3, 1, 1], 7: [7, 2, 3], 15: [15, 3, 7]}[conv1_size]
        self.conv1 = nn.Conv2d(3, 64, kernel_size=kernel_size, stride=stride, padding=padding, bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1, activation_fn=self.activation_fn)
        self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2, activation_fn=self.activation_fn)
        self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2, activation_fn=self.activation_fn)
        self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2, activation_fn=self.activation_fn)
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.linear = nn.Linear(512 * block.expansion, num_classes)

        self.normalize = None

    def _make_layer(self, block, planes, num_blocks, stride, activation_fn=nn.ReLU()):
        strides = [stride] + [1] * (num_blocks - 1)
        layers = []
        for stride in strides:
            layers.append(block(self.in_planes, planes, stride, activation_fn))
            self.in_planes = planes * block.expansion
        return nn.Sequential(*layers)

    def forward(self, x, penu=False):
        if not self.normalize:
            x = self.normalize(x)
        out = self.activation_fn(self.bn1(self.conv1(x)))
        out = self.layer1(out)
        out = self.layer2(out)
        out = self.layer3(out)
        out = self.layer4(out)
        out = self.avgpool(out)
        out = out.view(out.size(0), -1)
        if penu:
            return out
        out = self.linear(out)
        return out

def ResNet18(num_classes=10, conv1_size=3, activation_fn=nn.ReLU):
    return ResNet(BasicBlock, [2, 2, 2, 2], num_classes=num_classes, conv1_size=conv1_size, activation_fn=activation_fn)

"""# attack_util.py"""

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable

#dtype = torch.FloatTensor
dtype = torch.cuda.FloatTensor # run on GPU

### Do not modif the following codes
class ctx_noparamgrad(object):
    def __init__(self, module):
        self.prev_grad_state = get_param_grad_state(module)
        self.module = module
        set_param_grad_off(module)

    def __enter__(self):
        pass

    def __exit__(self, *args):
        set_param_grad_state(self.module, self.prev_grad_state)
        return False
        
def get_param_grad_state(module):
    return {param: param.requires_grad for param in module.parameters()}

def set_param_grad_off(module):
    for param in module.parameters():
        param.requires_grad = False

def set_param_grad_state(module, grad_state):
    for param in module.parameters():
        param.requires_grad = grad_state[param]

### Ends


### PGD Attack
class PGDAttack():
    def __init__(self, step = 10, eps = 8 / 255, alpha = 0.01, loss_type = 'ce', targeted = True, 
                    num_classes = 10, norm = 'linf', fgsm = False):
        '''
        norm: this parameter means which type of l-p norm constraints we use for attack. Note that we onlyuse L-inf norm for our homework. 
        Therefore, this parameter is always linf.
        But if you are interested in implementing an l-2 norm bounded attack, you can also try to implement it. Note that in that case,
        the eps should be set to a larger value such as 200/255 because of the difference between l-2 and l-inf.
        '''
        self.attack_step = step
        self.eps = eps
        self.alpha = alpha
        self.loss_type = loss_type
        self.targeted = targeted
        self.num_classes = num_classes
        self.norm = norm
        self.fgsm = fgsm


    def ce_loss(self, logits, ys, reduction = 'none'):
        ### Your code here
        return nn.CrossEntropyLoss()(logits, ys)
        ### Your code ends

    def cw_loss(self, logits, ys, reduction = 'none'):
        ### Your code here
        
        if self.targeted:
            target_class = 1
            target = torch.ones(logits.shape[0]).long()
            target = target.to(device)
            logits_wo_target = logits[(1 - torch.nn.functional.one_hot(target, logits.shape[1])).bool()].reshape(logits.shape[0], -1)  
            
            target_col = logits.gather(1, target.view(-1,1)).view(logits.shape[0]).type(dtype) # this is col 1
            target_col = target_col.to(device)
            max_p_without_target = torch.max(logits_wo_target,1).values.type(dtype) # this is max of all probas
            diff = torch.sub(max_p_without_target, target_col).type(dtype)
            all_zeros = torch.full((1, logits.shape[0]), 0.0).type(dtype)
            max_vect = torch.maximum(all_zeros, diff).type(dtype)
            return torch.mean(max_vect)

        else: # if untargeted
            proba_incorrect = logits[(1 - torch.nn.functional.one_hot(ys, logits.shape[1])).bool()].reshape(logits.shape[0], -1)
            correct_predict = logits.gather(1, ys.view(-1,1)).view(logits.shape[0]).type(dtype)
            max_proba_incorrect = torch.max(proba_incorrect,1).values.type(dtype) # this is max of all probas
            diff = torch.sub(correct_predict, max_proba_incorrect).type(dtype)
            all_zeros = torch.full((1, logits.shape[0]), 0.0).type(dtype)
            max_vect = torch.maximum(all_zeros, diff).type(dtype)
            return torch.mean(max_vect)

        ### Your code ends

    def clamp(self, delta, lower, upper):
        ### Your code here
        return torch.clamp(delta, lower, upper)
        ### Your code ends

    def linf_proj(self, delta):
        ### Your code here
        return torch.clamp(delta, -self.eps, self.eps)
        ### Your code ends

    def perturb(self, model: nn.Module, Xs, ys):
        #print ("PGD perturb called\n")
        delta = torch.zeros_like(Xs).to(Xs).uniform_(-self.eps, self.eps)
        delta.requires_grad = True

        if self.targeted:
          loss_fct = self.cw_loss
        else:
          if loss_type == 'ce':
              #print ("loss fct is ce\n")
              loss_fct = self.ce_loss
          elif loss_type == 'cw':
              loss_fct = self.cw_loss
        
        for iter_idx in range(self.attack_step):
            ### Your code here
            loss = loss_fct(model(Xs + delta), ys)
            # Calculate the gradients
            loss.backward()
            
            delta.data = self.linf_proj(delta + self.alpha*delta.grad.detach().sign())
            delta.grad.zero_()
        return delta.detach() 
        ### Your code ends  



### FGSMAttack
'''
Theoretically you can transform your PGDAttack to FGSM Attack by controling some of its parameters like `attack_step`. 
If you do that, you do not need to implement FGSM in this class.
'''
class FGSMAttack():
    def __init__(self, eps = 8 / 255, loss_type = 'ce', targeted = True, num_classes = 10, norm = 'linf'):
        self.eps = eps
        self.loss_type = loss_type
        self.targeted = targeted
        self.num_classes = num_classes
        self.norm = norm

    def perturb(self, model: nn.Module, Xs, ys):
        #print ("FGSM perturb called\n")
        delta = torch.zeros_like(Xs).to(Xs)
        delta.requires_grad = True
        ### Your code here
        if loss_type == 'ce':
            loss = PGDAttack.ce_loss(self, model(Xs + delta), ys)
        # Calculate the gradients
        loss.backward()        
        delta = self.eps * delta.grad.detach().sign()
        return delta      
        ### Your code ends

"""# evaluate.py"""

import torch
#import data_util
#import model_util
#import attack_util
#from attack_util import ctx_noparamgrad

from tqdm import tqdm
from io import BytesIO

eps = 6
alpha = 2
attack_rs = 1
attack_step = 10
loss_type = 'ce'
data_dir ='./data/'
model_prefix = './checkpoints/'
              
model_name ='resnet_cifar10.pth'
#model_name ='pgd10_eps8.pth'
fgsm = False
targeted = False


device = torch.device(0) if torch.cuda.is_available() else torch.device("cpu")
train_loader, valid_loader, test_loader, norm_layer = cifar10_dataloader(data_dir = data_dir)
num_classes = 10
model = ResNet18(num_classes = num_classes)
model.normalize = norm_layer
model_path = model_prefix + model_name

with open(model_path, 'rb') as fh:
    buf = BytesIO(fh.read())

model.load(buf) 
#model.load(model_path, map_location=torch.device('cpu'))
model = model.to(device)

#attack_step = attack_step
eps = eps / 255
restart = attack_rs
alpha = alpha / 255
#fgsm = fgsm
#loss_type = loss_type
#targeted = targeted

### Your code here for creating the attacker object
if fgsm:
    attacker = FGSMAttack(eps = eps, loss_type = loss_type, targeted = targeted, num_classes = num_classes, norm = 'linf')
else:
    attacker = PGDAttack(step = attack_step, eps = eps, alpha = alpha, 
                                    loss_type = loss_type, targeted = targeted, num_classes = num_classes, 
                                    norm = 'linf', fgsm = fgsm)

### Your code ends

total = 0
clean_correct_num = 0
robust_correct_num = 0
target_label = 1 ## only for targeted attack


## Make sure the model is in `eval` mode. Otherwise some operations such as dropout will  
model.eval()

for data, labels in tqdm(test_loader):
    data = data.float().to(device)
    if targeted:
        data_mask = (labels != target_label)
        if data_mask.sum() == 0:
            continue
        data = data[data_mask]
        labels = labels[data_mask]
        attack_labels = torch.ones_like(labels).to(device)
    else:
        attack_labels = labels
    attack_labels = attack_labels.to(device)
    labels = labels.to(device) # added from email
    batch_size = data.size(0)
    total += batch_size
    with ctx_noparamgrad(model):
        ### generate perturbation
        perturbed_data = attacker.perturb(model, data, attack_labels) + data

        predictions = model(data)
        clean_correct_num += torch.sum(torch.argmax(predictions, dim = -1) == labels).item()

        predictions = model(perturbed_data)
        robust_correct_num += torch.sum(torch.argmax(predictions, dim = -1) == labels).item()

print(f"total {total}, correct {clean_correct_num}, adversarial correct {robust_correct_num}, clean accuracy {clean_correct_num / total}, robust accuracy {robust_correct_num / total}")