-
Notifications
You must be signed in to change notification settings - Fork 7
/
loss_functions.py
111 lines (85 loc) · 3.77 KB
/
loss_functions.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
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from util_MNIST import retrieveMNISTTrainingData
from util_model import MNISTClassifier, trainModel
from adversarial_attack_DRO import ProjetcedDRO, LagrangianDRO, FrankWolfeDRO
"""
This module contains the seven loss functions listed in Carlini and Wagneer
(2017) with a minor modification in the constant in f_5.
"""
def f_1(outputs, labels):
return F.cross_entropy(outputs, labels)
def f_2(outputs, labels):
outputs = F.softmax(outputs, dim=1)
return f_6(outputs, labels)
def f_3(outputs, labels):
outputs = F.softmax(outputs, dim=1)
return f_7(outputs, labels)
def f_4(outputs, labels):
outputs = F.softmax(outputs, dim=1)
reference_outputs = torch.gather(
outputs, 1, labels.view(-1, 1).long()).view(-1)
return torch.mean(torch.clamp(0.5 - reference_outputs, min=0))
def f_5(outputs, labels):
# Note that in the original version, the base of e is used instead of 2.
outputs = F.softmax(outputs, dim=1)
reference_outputs = torch.gather(
outputs, 1, labels.view(-1, 1).long()).view(-1)
return torch.mean(torch.log2(2.125 - 2 * reference_outputs))
def f_6(outputs, labels):
max_outputs, _ = torch.max(outputs, dim=1)
reference_outputs = torch.gather(
outputs, 1, labels.view(-1, 1).long()).view(-1)
return torch.mean(max_outputs - reference_outputs)
def f_7(outputs, labels, nb_classes=10):
#device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
device = "cuda:0"
batch_size = labels.size()[0]
indexes_row = torch.arange(0, nb_classes).to(device)
indexes = indexes_row.repeat(batch_size, 1)
labels_cloned = labels.view(-1, 1).repeat(1, nb_classes)
new_outputs = outputs[indexes != labels_cloned].view(
batch_size, nb_classes-1)
reference_outputs = torch.gather(
outputs, 1, labels.view(-1, 1).long()).view(-1)
difference = torch.max(new_outputs, dim=1)[0] - reference_outputs
return torch.mean(F.softplus(difference))
def trainModelLoss(dro_type, epochs, steps_adv, budget, activation, batch_size, loss_criterion, cost_function=None):
"""
Train a neural network with a specified loss function.
"""
model = MNISTClassifier(activation=activation)
if dro_type == 'PGD':
train_module = ProjetcedDRO(model, loss_criterion)
elif dro_type == 'Lag':
assert cost_function is not None
train_module = LagrangianDRO(model, loss_criterion, cost_function)
elif dro_type == 'FW':
train_module = FrankWolfeDRO(model, loss_criterion, p=2, q=2)
else:
raise ValueError("The type of DRO is not valid.")
train_module.train(budget=budget, batch_size=batch_size,
epochs=epochs, steps_adv=steps_adv)
folderpath = "./Loss_models/"
filepath = folderpath + "{}_DRO_activation={}_epsilon={}_loss={}.pt".format(
dro_type, activation, budget, loss_criterion.__name__)
torch.save(model.state_dict(), filepath)
print("A neural network adversarially trained using {} now saved at: {}".format(
dro_type, filepath))
if __name__ == "__main__":
epochs = 25
steps_adv = 15
epsilon = 2.8
optimal_gamma = 1.0
batch_size = 128
loss_criterions = [f_1, f_2, f_3, f_4, f_5, f_6, f_7]
def cost_function(x, y): return torch.dist(x, y, p=2) ** 2
for loss_criterion in loss_criterions:
trainModelLoss("FW", epochs, steps_adv, epsilon,
"relu", batch_size, loss_criterion)
trainModelLoss("PGD", epochs, steps_adv, epsilon,
"relu", batch_size, loss_criterion)
trainModelLoss("Lag", epochs, steps_adv, optimal_gamma, "relu",
batch_size, loss_criterion, cost_function=cost_function)