-
Notifications
You must be signed in to change notification settings - Fork 10
/
utils.py
47 lines (40 loc) · 1.79 KB
/
utils.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
import torch
import numpy as np
"""Normalize the data given the dataset. Only ImageNet and CIFAR-10 are supported"""
def transform(img, dataset='imagenet'):
# Data
if dataset == 'imagenet':
mean = torch.Tensor([0.485, 0.456, 0.406]).unsqueeze(1).expand_as(img[0, :, :, 0]).unsqueeze(2).expand_as(
img[0]).unsqueeze(0).expand_as(img).cuda()
std = torch.Tensor([0.229, 0.224, 0.225]).unsqueeze(1).expand_as(img[0, :, :, 0]).unsqueeze(2).expand_as(
img[0]).unsqueeze(0).expand_as(img).cuda()
elif dataset == 'cifar':
mean = torch.Tensor([0.485, 0.456, 0.406]).unsqueeze(1).expand_as(img[0, :, :, 0]).unsqueeze(2).expand_as(
img[0]).unsqueeze(0).expand_as(img).cuda()
std = torch.Tensor([0.229, 0.224, 0.225]).unsqueeze(1).expand_as(img[0, :, :, 0]).unsqueeze(2).expand_as(
img[0]).unsqueeze(0).expand_as(img).cuda()
else:
raise "dataset is not supported"
return (img - mean) / std
"""Given [label] and [dataset], return a random label different from [label]"""
def random_label(label, dataset='imagenet'):
if dataset == 'imagenet':
class_num = 1000
elif dataset == 'cifar':
class_num = 10
else:
raise "dataset is not supported"
attack_label = np.random.randint(class_num)
while label == attack_label:
attack_label = np.random.randint(class_num)
return attack_label
"""Given the variance of zero_mean Gaussian [n_radius], return a noisy version of [img]"""
def noisy_img(img, n_radius):
return img + n_radius * torch.randn_like(img)
class Noisy(torch.autograd.Function):
@staticmethod
def forward(self, img, n_radius):
return noisy_img(img, n_radius=n_radius)
@staticmethod
def backward(self, grad_output):
return grad_output, None