File size: 4,883 Bytes
e1aaaac |
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 |
import tensorflow as tf
import numpy as np
import torch
class ModelAdapter():
def __init__(self, logits, x, y, sess, num_classes=10):
self.logits = logits
self.sess = sess
self.x_input = x
self.y_input = y
self.num_classes = num_classes
# gradients of logits
if num_classes <= 10:
self.grads = [None] * num_classes
for cl in range(num_classes):
self.grads[cl] = tf.gradients(self.logits[:, cl], self.x_input)[0]
# cross-entropy loss
self.xent = tf.nn.sparse_softmax_cross_entropy_with_logits(
logits=self.logits, labels=self.y_input)
self.grad_xent = tf.gradients(self.xent, self.x_input)[0]
# dlr loss
self.dlr = dlr_loss(self.logits, self.y_input, num_classes=self.num_classes)
self.grad_dlr = tf.gradients(self.dlr, self.x_input)[0]
# targeted dlr loss
self.y_target = tf.placeholder(tf.int64, shape=[None])
self.dlr_target = dlr_loss_targeted(self.logits, self.y_input, self.y_target, num_classes=self.num_classes)
self.grad_target = tf.gradients(self.dlr_target, self.x_input)[0]
self.la = tf.placeholder(tf.int64, shape=[None])
self.la_target = tf.placeholder(tf.int64, shape=[None])
la_mask = tf.one_hot(self.la, self.num_classes)
la_target_mask = tf.one_hot(self.la_target, self.num_classes)
la_logit = tf.reduce_sum(la_mask * self.logits, axis=1)
la_target_logit = tf.reduce_sum(la_target_mask * self.logits, axis=1)
self.diff_logits = la_target_logit - la_logit
self.grad_diff_logits = tf.gradients(self.diff_logits, self.x_input)[0]
def predict(self, x):
x2 = np.moveaxis(x.cpu().numpy(), 1, 3)
y = self.sess.run(self.logits, {self.x_input: x2})
return torch.from_numpy(y).cuda()
def grad_logits(self, x):
x2 = np.moveaxis(x.cpu().numpy(), 1, 3)
logits, g2 = self.sess.run([self.logits, self.grads], {self.x_input: x2})
g2 = np.moveaxis(np.array(g2), 0, 1)
g2 = np.transpose(g2, (0, 1, 4, 2, 3))
return torch.from_numpy(logits).cuda(), torch.from_numpy(g2).cuda()
def get_grad_diff_logits_target(self, x, y=None, y_target=None):
la = y.cpu().numpy()
la_target = y_target.cpu().numpy()
x2 = np.moveaxis(x.cpu().numpy(), 1, 3)
dl, g2 = self.sess.run([self.diff_logits, self.grad_diff_logits], {self.x_input: x2, self.la: la, self.la_target: la_target})
g2 = np.transpose(np.array(g2), (0, 3, 1, 2))
return torch.from_numpy(dl).cuda(), torch.from_numpy(g2).cuda()
def get_logits_loss_grad_xent(self, x, y):
x2 = np.moveaxis(x.cpu().numpy(), 1, 3)
y2 = y.clone().cpu().numpy()
logits_val, loss_indiv_val, grad_val = self.sess.run([self.logits, self.xent, self.grad_xent], {self.x_input: x2, self.y_input: y2})
grad_val = np.moveaxis(grad_val, 3, 1)
return torch.from_numpy(logits_val).cuda(), torch.from_numpy(loss_indiv_val).cuda(), torch.from_numpy(grad_val).cuda()
def get_logits_loss_grad_dlr(self, x, y):
x2 = np.moveaxis(x.cpu().numpy(), 1, 3)
y2 = y.clone().cpu().numpy()
logits_val, loss_indiv_val, grad_val = self.sess.run([self.logits, self.dlr, self.grad_dlr], {self.x_input: x2, self.y_input: y2})
grad_val = np.moveaxis(grad_val, 3, 1)
return torch.from_numpy(logits_val).cuda(), torch.from_numpy(loss_indiv_val).cuda(), torch.from_numpy(grad_val).cuda()
def get_logits_loss_grad_target(self, x, y, y_target):
x2 = np.moveaxis(x.cpu().numpy(), 1, 3)
y2 = y.clone().cpu().numpy()
y_targ = y_target.clone().cpu().numpy()
logits_val, loss_indiv_val, grad_val = self.sess.run([self.logits, self.dlr_target, self.grad_target], {self.x_input: x2, self.y_input: y2, self.y_target: y_targ})
grad_val = np.moveaxis(grad_val, 3, 1)
return torch.from_numpy(logits_val).cuda(), torch.from_numpy(loss_indiv_val).cuda(), torch.from_numpy(grad_val).cuda()
def dlr_loss(x, y, num_classes=10):
x_sort = tf.contrib.framework.sort(x, axis=1)
y_onehot = tf.one_hot(y, num_classes)
### TODO: adapt to the case when the point is already misclassified
loss = -(x_sort[:, -1] - x_sort[:, -2]) / (x_sort[:, -1] - x_sort[:, -3] + 1e-12)
return loss
def dlr_loss_targeted(x, y, y_target, num_classes=10):
x_sort = tf.contrib.framework.sort(x, axis=1)
y_onehot = tf.one_hot(y, num_classes)
y_target_onehot = tf.one_hot(y_target, num_classes)
loss = -(tf.reduce_sum(x * y_onehot, axis=1) - tf.reduce_sum(x * y_target_onehot, axis=1)) / (x_sort[:, -1] - .5 * x_sort[:, -3] - .5 * x_sort[:, -4] + 1e-12)
return loss
|