Spaces:
Runtime error
Runtime error
| # Helper function for extracting features from pre-trained models | |
| import os | |
| import sys | |
| import math | |
| import numbers | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import torchvision.transforms as transforms | |
| from torch.autograd import Variable | |
| import torch.nn as nn | |
| from PIL import Image | |
| import numpy as np | |
| from util.feature_extraction_utils import warp_image, normalize_batch | |
| from util.prepare_utils import get_ensemble, extract_features | |
| from lpips_pytorch import LPIPS, lpips | |
| from tqdm import tqdm | |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| tensor_transform = transforms.ToTensor() | |
| pil_transform = transforms.ToPILImage() | |
| class Attack(nn.Module): | |
| def __init__(self, models, dim, attack_type, eps, c_sim=0.5, net_type='alex', lr=0.05, | |
| n_iters=100, noise_size=0.001, n_starts=10, c_tv=None, sigma_gf=None, kernel_size_gf=None, | |
| combination=False, warp=False, theta_warp=None, V_reduction=None): | |
| super(Attack, self).__init__() | |
| self.extractor_ens = get_ensemble(models, sigma_gf, kernel_size_gf, combination, V_reduction, warp, theta_warp) | |
| #print("There are '{}'' models in the attack ensemble".format(len(self.extractor_ens))) | |
| self.dim = dim | |
| self.eps = eps | |
| self.c_sim = c_sim | |
| self.net_type = net_type | |
| self.lr = lr | |
| self.n_iters = n_iters | |
| self.noise_size = noise_size | |
| self.n_starts = n_starts | |
| self.c_tv=None | |
| self.attack_type = attack_type | |
| self.warp = warp | |
| self.theta_warp = theta_warp | |
| if self.attack_type == 'lpips': | |
| self.lpips_loss = LPIPS(self.net_type).to(device) | |
| def execute(self, images, dir_vec, direction): | |
| images = Variable(images).to(device) | |
| dir_vec = dir_vec.to(device) | |
| # take norm wrt dim | |
| dir_vec_norm = dir_vec.norm(dim = 2).unsqueeze(2).to(device) | |
| dist = torch.zeros(images.shape[0]).to(device) | |
| adv_images = images.detach().clone() | |
| if self.warp: | |
| self.face_img = warp_image(images, self.theta_warp) | |
| for start in range(self.n_starts): | |
| # update adversarial images old and distance old | |
| adv_images_old = adv_images.detach().clone() | |
| dist_old = dist.clone() | |
| # add noise to initialize ( - noise_size, noise_size) | |
| noise_uniform = Variable(2 * self.noise_size * torch.rand(images.size()) - self.noise_size).to(device) | |
| adv_images = Variable(images.detach().clone() + noise_uniform, requires_grad=True).to(device) | |
| for i in range(self.n_iters): | |
| adv_features = extract_features(adv_images, self.extractor_ens, self.dim).to(device) | |
| # normalize feature vectors in ensembles | |
| loss = direction*torch.mean((adv_features - dir_vec)**2/dir_vec_norm) | |
| if self.c_tv != None: | |
| tv_out = self.total_var_reg(images, adv_images) | |
| loss -= self.c_tv * tv_out | |
| if self.attack_type == 'lpips': | |
| lpips_out = self.lpips_reg(images, adv_images) | |
| loss -= self.c_sim * lpips_out | |
| grad = torch.autograd.grad(loss, [adv_images]) | |
| adv_images = adv_images + self.lr * grad[0].sign() | |
| perturbation = adv_images - images | |
| if self.attack_type == 'sgd': | |
| perturbation = torch.clamp(perturbation, min=-self.eps, max=self.eps) | |
| adv_images = images + perturbation | |
| adv_images = torch.clamp(adv_images, min=0, max=1) | |
| adv_features = extract_features(adv_images, self.extractor_ens, self.dim).to(device) | |
| dist = torch.mean((adv_features - dir_vec)**2/dir_vec_norm, dim = [1,2]) | |
| if direction ==1: | |
| adv_images[dist < dist_old] = adv_images_old[dist < dist_old] | |
| dist[dist < dist_old] = dist_old[dist < dist_old] | |
| else: | |
| adv_images[dist > dist_old] = adv_images_old[dist > dist_old] | |
| dist[dist > dist_old] = dist_old[dist > dist_old] | |
| return adv_images.detach().cpu() | |
| def lpips_reg(self, images, adv_images): | |
| if self.warp: | |
| face_adv = warp_image(adv_images, self.theta_warp) | |
| lpips_out = self.lpips_loss(normalize_batch(self.face_img).to(device), normalize_batch(face_adv).to(device))[0][0][0][0] /(2*adv_images.shape[0]) | |
| lpips_out += self.lpips_loss(normalize_batch(images).to(device), normalize_batch(adv_images).to(device))[0][0][0][0] / (2*adv_images.shape[0]) | |
| else: | |
| lpips_out = self.lpips_loss(normalize_batch(images).to(device), normalize_batch(adv_images).to(device))[0][0][0][0] / adv_images.shape[0] | |
| return lpips_out | |
| def total_var_reg(images, adv_images): | |
| perturbation = adv_images - images | |
| tv = torch.mean(torch.abs(perturbation[:, :, :, :-1] - perturbation[:, :, :, 1:])) + \ | |
| torch.mean(torch.abs(perturbation[:, :, :-1, :] - perturbation[:, :, 1:, :])) | |
| return tv | |