# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: pca_gm
#     language: python
#     name: pca_gm
# ---

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"

# +
import itertools
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import time
import torchvision
import torch.optim as optim
from torchvision import datasets, transforms
from pgd_attack import pgd_attack, margin_loss, TransferPGD
from models import ResNet18, WideResNet28, ResNet34, WideResNet34
from tqdm import tqdm, trange
from model import get_model_for_attack

from attack_main import eval_model_pgd, eval_model_with_attack
from utils import prepare_cifar, Logger, check_mkdir
from eval_model import eval_model, eval_model_pgd
# -

import matplotlib.pyplot as plt

torch.manual_seed(0)


def get_train_data():
    trainset_org = torchvision.datasets.CIFAR10(
        root='./data', train=True, download=True,
        transform=transforms.ToTensor()
    )
    loader_org = torch.utils.data.DataLoader(trainset_org, 10)
    return itertools.chain(loader_org)


dataset = get_train_data()

model = WideResNet34()
model.load_state_dict(torch.load('logs/Jun08-2018_WideResNet34/WideResNet34_e49_0.8695_0.5897-final.pt'))
model.eval()

modeln = get_model_for_attack('model1')

eg = None
egy = None
for x, y in dataset:
    eg = x.squeeze()
    egy = y
    break

egy.shape

plt.imshow(eg[0].permute((1,2,0)))

print("true val: ", egy)
print('predict raw', torch.argmax(modeln(eg),dim=1) )
for i in range(1):
    x_adv = eg + 0.01 * torch.randn(eg.shape).detach()
    x_adv = torch.clamp(x_adv, 0.0, 1.0)
    plt.imshow(x_adv[0].permute(1,2,0))
    logits = modeln(x_adv) 
    print('pred: ', torch.argmax(logits, dim=1), 'logit: ', logits)

print("true val: ", egy)
print('predict raw', torch.argmax(modeln(eg),dim=1) )
for i in range(1):
    x_adv = pgd_attack(modeln, eg.clone(), egy.clone(), step_size=0.01, epsilon=8.0/255, perturb_steps=20)
    plt.imshow(x_adv[0].permute(1,2,0))
    logits = modeln(x_adv) 
    print('pred: ', torch.argmax(logits, dim=1), 'logit: ', logits)

# ## Robust Model

print("true val: ", egy)
print('predict raw', torch.argmax(model(eg),dim=1) )
for i in range(1):
    x_adv = eg + 0.01 * torch.randn(eg.shape).detach()
    x_adv = torch.clamp(x_adv, 0.0, 1.0)
    plt.imshow(x_adv[0].permute(1,2,0))
    logits = model(x_adv) 
    print('pred: ', torch.argmax(logits, dim=1), 'logit: ', logits)

print("true val: ", egy)
print('predict raw', torch.argmax(model(eg),dim=1) )
for i in range(1):
    x_adv = pgd_attack(model, eg.clone(), egy.clone(), step_size=0.01, epsilon=8.0/255, perturb_steps=20)
    plt.imshow(x_adv[0].permute(1,2,0))
    logits = model(x_adv) 
    print('pred: ', torch.argmax(logits, dim=1), 'logit: ', logits)

model4 = get_model_for_attack('model4')

logits4 = model4(x_adv) 
print('pred: ', torch.argmax(logits4, dim=1), 'logit: ', logits)

# ## Guessing the variance

import numpy as np

# +
norm_n = []
norm_r6 = []
norm_r4 = []

model = model.cpu()
modeln = modeln.cpu()
model4 = model4.cpu()

trainset_org = torchvision.datasets.CIFAR10(
        root='./data', train=True, download=True,
        transform=transforms.ToTensor()
    )
loader_org = torch.utils.data.DataLoader(trainset_org, 50)

for i, (x, y) in tqdm(enumerate(loader_org)):
    logits_n = modeln(x)
    logits_r = model(x)
    logits_4 = model4(x)
    norm_n += [*logits_n.norm(dim=1).detach().numpy()]
    norm_r6 += [*logits_r.norm(dim=1).detach().numpy()]
    norm_r4 += [*logits_4.norm(dim=1).detach().numpy()]
    if (i == 30): break
# -

plt.subplot(2,1,1)
plt.hist(norm_n, bins=30)
plt.subplot(2,1,2)
plt.hist(norm_r, bins=30)

plt.style.use('seaborn-pastel')
plt.hist(norm_n, bins=30, label='Naive model')
plt.hist(norm_r6, bins=30, label='robust model6')
plt.hist(norm_r4, bins=30, label='robust model4')
plt.legend()
plt.xlabel('Variance of logits')
plt.ylabel('Frequency')
plt.savefig('variance.png')

# +
loss_n = []
loss_r6 = []
loss_r4 = []

trainset_org = torchvision.datasets.CIFAR10(
        root='./data', train=True, download=True,
        transform=transforms.ToTensor()
    )
loader_org = torch.utils.data.DataLoader(trainset_org, 10)

for i, (x, y) in tqdm(enumerate(loader_org)):
    logits_n = nn.CrossEntropyLoss()(modeln(x),y)
    logits_r = nn.CrossEntropyLoss()(model(x),y)
    logits_4 = nn.CrossEntropyLoss()(model4(x),y)
    loss_n += [logits_n.detach().numpy()]
    loss_r6 += [logits_r.detach().numpy()]
    loss_r4 += [logits_4.detach().numpy()]
    if (i == 30): break
# -

plt.style.use('seaborn-pastel')
plt.hist(loss_n, bins=30, label='Naive model')
plt.hist(loss_r6, bins=30, label='robust model6')
plt.hist(loss_r4, bins=30, label='robust model4')
plt.legend()
plt.xlabel('Variance of logits')
plt.ylabel('Frequency')
plt.savefig('variance.png')

# +
## Explore random space

# +
loss_n = []
loss_r6 = []
loss_r4 = []

for i in tqdm(range(500)):
    x = eg + 0.01 * torch.randn(eg.shape).detach()
    x = torch.clamp(x, 0.0, 1.0)
    logits_n = nn.CrossEntropyLoss()(modeln(x),egy)
    logits_r = nn.CrossEntropyLoss()(model(x),egy)
    logits_4 = nn.CrossEntropyLoss()(model4(x),egy)
    loss_n += [logits_n.detach().numpy()]
    loss_r6 += [logits_r.detach().numpy()]
    loss_r4 += [logits_4.detach().numpy()]

# -

plt.style.use('seaborn-pastel')
plt.hist(loss_n,  label='Naive model')
plt.hist(loss_r6,  label='robust model6')
plt.hist(loss_r4,  label='robust model4')
plt.legend()
plt.xlabel('Variance of logits')
plt.ylabel('Frequency')
plt.savefig('variance.png')

# ## Fail

# +
norm_adv_n = []
norm_adv_r = []
acc = 0 
racc = 0

trainset_org = torchvision.datasets.CIFAR10(
        root='./data', train=True, download=True,
        transform=transforms.ToTensor()
    )
loader_org = torch.utils.data.DataLoader(trainset_org, 256)

for i, (x, y) in tqdm(enumerate(loader_org)):
    x = x.cuda()
    y = y.cuda()
    modeln.cuda()
    model.cuda()
    x_adv = pgd_attack(model, x.clone(), y.clone(), step_size=0.01, epsilon=8.0/255, perturb_steps=20)
    logits_n = modeln(x_adv)
    logits_r = model(x_adv)
    norm_adv_n += [*logits_n.norm(dim=1).detach().cpu().numpy()]
    norm_adv_r += [*logits_r.norm(dim=1).detach().cpu().numpy()]
    acc += (torch.argmax(logits_n,dim=1) == y).sum()
    racc +=(torch.argmax(logits_r,dim=1) == y).sum()
# -

print('acc= {}, racc= {}'.format(acc,racc))
plt.subplot(2,1,1)
plt.hist(norm_n, bins=30)
plt.subplot(2,1,2)
plt.hist(norm_r, bins=30)



StarcoderdataJupyterScriptsDedupFiltered