-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathprotector.py
153 lines (113 loc) · 3.98 KB
/
protector.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
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
import torch
import numpy as np
import pandas as pd
from math import log, sqrt
import pickle
import matplotlib.pyplot as plt
from loguru import logger
import copy
from cdf import CDF
# from exp_utils.utils import get_ents_acc
class Protector:
def __init__(self, cdf, device, **kwargs):
self.cdf = cdf
self.device = device
self.C = 1.
self.D = kwargs.get('eps_clip_val', 1.80)
self.gamma = kwargs.get('gamma', 2 / sqrt(3))
self.martingales = []
self.gradients = []
self.epsilons = [0]
self.info = {
'u_before': [],
'u_after': [],
'z_before': [],
'z_after': [],
'last_batch_u_before': [],
'last_batch_u_after': [],
}
@property
def last_eps(self):
return self.epsilons[-1]
def set_eps_clip_val(self, eps_clip_val):
logger.info(f"setting epsilon clip val to {eps_clip_val}")
self.D = eps_clip_val
def set_gamma(self, gamma):
logger.info(f"setting gamma val to {gamma}")
self.gamma = gamma
def export_info(self):
return {
'u_before': np.array(self.info['u_before']).reshape(-1, 1),
'u_after': np.array(self.info['u_after']).reshape(-1, 1),
'z_before': np.array(self.info['z_before']).reshape(-1, 1),
'z_after': np.array(self.info['z_after']).reshape(-1, 1),
'martingales': np.array(self.martingales).reshape(-1, 1),
'epsilons': np.array(self.epsilons).reshape(-1, 1),
}
def reset(self):
self.C = 1.
self.martingales = []
self.gradients = []
self.epsilons = [0]
self.info = {
'u_before': [],
'u_after': [],
'z_before': [],
'z_after': [],
'last_batch_u_before': [],
'last_batch_u_after': [],
}
def protect(self, z):
self.info['last_batch_u_before'] = []
self.info['last_batch_u_after'] = []
zz = z.clone().cpu().detach().numpy()
for j in range(zz.shape[0]):
self.info['z_before'].append(float(zz[j]))
u = self.cdf(zz[j])
u_protected = self.protect_u(u)
z_fixed = self.cdf.inverse(u_protected)
zz[j] = z_fixed
self.info['u_before'].append(float(u))
self.info['u_after'].append(float(u_protected))
self.info['z_after'].append(float(z_fixed))
self.info['last_batch_u_before'].append(float(u))
self.info['last_batch_u_after'].append(float(u_protected))
zz = torch.tensor(zz).to(self.device)
return zz, self.info
def sfogd(self, u_t):
if len(self.epsilons) == 0:
return 0
if len(self.info['u_before']) == 0:
return 0
eps_t = self.last_eps
v_t = u_t - 0.5
E_tau = self.D * np.sign(u_t - 0.50)
ind = 0 if E_tau * eps_t > 0 and np.abs(eps_t) > self.D else 1
grad_t = (v_t / (1 + eps_t * v_t)) * ind
self.gradients.append(grad_t)
if grad_t != 0:
grad_arr = np.array(self.gradients)
c = self.gamma * ((grad_t) / (np.sqrt((grad_arr**2).sum())))
eps_new = eps_t + c
else:
eps_new = eps_t
return eps_new
def protect_u(self, u_t):
C = self.C
eps_t = self.last_eps
b = 1 + eps_t * (u_t - 0.5)
u_protected = 0.5 * eps_t * (u_t ** 2) + (1 - 0.5 * eps_t) * u_t
eps_new = self.sfogd(u_t)
self.C = min(float(C * b), 1e200)
self.martingales.append(self.C)
self.epsilons.append(float(eps_new))
return u_protected
def get_protector_from_ents(ents, args):
logger.info('creating protector from ents')
cdf = CDF(ents)
gamma = args.gamma
eps_clip_val = args.eps_clip
protector = Protector(cdf, args.device)
protector.set_gamma(gamma)
protector.set_eps_clip_val(eps_clip_val)
return protector