def testVirtualAdvRegularizer(self):
"""Tests virtual_adv_regularizer returning expected loss."""
np_input = np.array([[1.0, -1.0]])
tf_input = tf.constant(np_input)
np_weights = np.array([[1.0, 5.0], [2.0, 2.0]])
tf_weights = tf.constant(np_weights)
# Linear transformation and L2 loss makes the Hessian matrix constant.
embedding_fn = lambda x: tf.matmul(x, tf_weights)
step_size = 0.1
vadv_config = configs.VirtualAdvConfig(
adv_neighbor_config=configs.AdvNeighborConfig(
feature_mask=None,
adv_step_size=step_size,
adv_grad_norm=configs.NormType.L2),
distance_config=configs.DistanceConfig(
distance_type=configs.DistanceType.L2, sum_over_axis=-1),
num_approx_steps=1,
approx_difference=1e-3) # enlarged for numerical stability
np_seed = np.array([[0.6, 0.8]])
tf_seed = tf.constant(np_seed)
vadv_loss = regularizer._virtual_adv_regularizer(
tf_input, embedding_fn, vadv_config, embedding_fn(tf_input),
tf_seed)
actual_loss = self.evaluate(vadv_loss)
# For detail derivation of the Hessian matrix, see go/vadv-tests-hessian
hessian = 2 * np.dot(np_weights, np_weights.T)
approx = np.matmul(np_seed, hessian)
approx *= step_size / np.linalg.norm(approx, axis=-1, keepdims=True)
expected_loss = np.linalg.norm(np.matmul(approx, np_weights))**2
self.assertNear(actual_loss, expected_loss, err=1e-5)
def _test_multi_iter_gen_adv_neighbor_should_ignore_sparse_tensors_setup(
self):
@tf.function
def model_fn(inp):
prod_sparse = tf.reduce_sum(
tf.sparse.sparse_dense_matmul(inp['sparse'], w_sparse))
prod_dense = tf.reduce_sum(input_tensor=w_dense * inp['dense'])
return prod_sparse + prod_dense
@tf.function
def loss_fn(label, pred):
return tf.abs(label - pred)
# sparse_feature represents [[1, 0, 2], [0, 3, 0]].
sparse_feature = tf.SparseTensor(
indices=[[0, 0], [0, 2], [1, 1]],
values=[1.0, 2.0, 3.0],
dense_shape=(2, 3))
x = {'sparse': sparse_feature, 'dense': tf.constant([[-1.0], [1.0]])}
w_sparse = tf.constant([[5.0], [4.0], [3.0]])
w_dense = tf.constant([[6.0]])
adv_config = configs.AdvNeighborConfig(
feature_mask=None,
adv_step_size=0.1,
adv_grad_norm='l2',
iterations=2,
epsilon=0.13)
return x, w_sparse, w_dense, adv_config, model_fn, loss_fn
def test_gen_adv_neighbor_for_feature_columns_with_int_feature_v2(self):
x, y, w, expected_neighbor_fc1, expected_neighbor_fc2 = (
self._test_gen_adv_neighbor_for_feature_columns_with_int_feature())
with tf.GradientTape() as tape:
tape.watch(x)
# Simple linear regression.
x_stacked = tf.concat(
[x['fc1'], x['fc2'],
tf.cast(x['fc_int'], tf.dtypes.float32)],
axis=1)
y_hat = tf.matmul(x_stacked, w)
loss = tf.math.squared_difference(y, y_hat)
adv_config = configs.AdvNeighborConfig(feature_mask=None,
adv_step_size=0.1,
adv_grad_norm='l2')
adv_neighbor, _ = adv_lib.gen_adv_neighbor(x,
loss,
adv_config,
gradient_tape=tape)
actual_neighbor = self.evaluate(adv_neighbor)
self.assertAllClose(expected_neighbor_fc1, actual_neighbor['fc1'])
self.assertAllClose(expected_neighbor_fc2, actual_neighbor['fc2'])
self.assertAllClose([[2]], actual_neighbor['fc_int'])
def _gen_adv_neighbor(x):
w = tf.constant([[3.0], [4.0]])
loss = tf.matmul(x, w)
adv_config = configs.AdvNeighborConfig(
feature_mask=None, adv_step_size=0.1, adv_grad_norm='l2')
adv_neighbor, _ = adv_lib.gen_adv_neighbor(x, loss, adv_config)
return adv_neighbor
def test_gen_adv_neighbor_multi_iter_respects_feature_constraints(
self, gen_adv_neighbor_fn):
w1 = tf.constant(1.0, shape=(2, 2, 3))
w2 = tf.constant(-1.0, shape=(2, 2))
f1 = np.array([[[[0.0, 0.1, 0.2], [0.3, 0.4, 0.5]],
[[0.6, 0.7, 0.8], [0.9, 1.0, 1.1]]]],
dtype=np.float32)
f2 = np.array([[[0.0, 0.33], [0.66, 1.0]]], dtype=np.float32)
x = {'f1': tf.constant(f1), 'f2': tf.constant(f2)}
y = tf.constant([0.0])
def model_fn(x):
return tf.reduce_sum(w1 * x['f1']) + tf.reduce_sum(w2 * x['f2'])
def loss_fn(_, pred):
return pred
adv_step_size = 0.5
adv_config = configs.AdvNeighborConfig(
adv_step_size=adv_step_size,
adv_grad_norm='infinity',
epsilon=0.4,
iterations=2,
clip_value_min={'f2': 0.0},
clip_value_max={'f1': 1.0})
adv_neighbor = gen_adv_neighbor_fn(x, y, model_fn, loss_fn, adv_config)
actual_neighbor = self.evaluate(adv_neighbor)
# gradient = w, perturbation = min(adv_step_size * sign(w), 0.4)
expected_neighbor = {
'f1': np.minimum(f1 + 0.4, 1.0),
'f2': np.maximum(f2 - 0.4, 0.0),
}
self.assertAllClose(expected_neighbor, actual_neighbor)
def test_multi_iter_gen_adv_neighbor_for_tensor_list(self):
x = [tf.constant([[-1.0]]), tf.constant([[1.0]])]
y = tf.constant([0.0])
w = [tf.constant([[3.0]]), tf.constant([[4.0]])]
loss_fn = tf.math.squared_difference
model_fn = (lambda inp: tf.matmul(inp[0], w[0]) + tf.matmul(inp[1], w[1]))
with tf.GradientTape() as tape:
tape.watch(x)
y_hat = tf.matmul(x[0], w[0]) + tf.matmul(x[1], w[1])
loss = tf.math.squared_difference(y, y_hat)
adv_config = configs.AdvNeighborConfig(
feature_mask=tf.constant(1.0),
adv_step_size=0.1,
adv_grad_norm='l2',
iterations=2)
adv_neighbor, _ = adv_lib.gen_adv_neighbor(
x,
loss,
adv_config,
gradient_tape=tape,
pgd_model_fn=model_fn,
pgd_loss_fn=loss_fn,
pgd_labels=y)
# gradient = [[6, 8]], normalized gradient = [[0.6, 0.8]]
expected_neighbor = [[[-1.0 + 0.1 * 0.6 * 2]], [[1.0 + 0.1 * 0.8 * 2]]]
actual_neighbor = self.evaluate(adv_neighbor)
self.assertAllClose(expected_neighbor, actual_neighbor)
def test_multi_iter_gen_adv_neighbor_proj_limits(self):
# Simple linear regression.
x = tf.constant([[-1.0, 1.0]])
y = tf.constant([0.0])
w = tf.constant([[3.0], [4.0]])
loss_fn = tf.math.squared_difference
model_fn = lambda input: tf.matmul(input, w)
with tf.GradientTape() as tape:
tape.watch(x)
y_hat = model_fn(x)
loss = loss_fn(y, y_hat)
adv_config = configs.AdvNeighborConfig(
feature_mask=tf.constant(1.0),
adv_step_size=0.1,
adv_grad_norm='l2',
iterations=2,
epsilon=0.15)
adv_neighbor, _ = adv_lib.gen_adv_neighbor(
x,
loss,
adv_config,
gradient_tape=tape,
pgd_model_fn=model_fn,
pgd_loss_fn=loss_fn,
pgd_labels=y)
# Take two steps in the gradient direction. Project back onto epsilon ball
# after iteration 2.
expected_neighbor = [[-1.0 + 0.1 * 0.6 * 1.5, 1.0 + 0.1 * 0.8 * 1.5]]
actual_neighbor = self.evaluate(adv_neighbor)
self.assertAllClose(expected_neighbor, actual_neighbor)
def _test_multi_iter_gen_adv_neighbor_for_features_with_different_shapes_setup(
self):
w1 = tf.constant(1.0, shape=(2, 2, 3))
w2 = tf.constant(1.0, shape=(2, 2))
f1 = tf.constant([[[[0.0, 0.1, 0.2], [0.3, 0.4, 0.5]],
[[0.6, 0.7, 0.8], [0.9, 1.0, 1.1]]],
[[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]],
[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]]])
f2 = tf.constant([[[1.2, 1.3], [1.4, 1.5]], [[0.0, 0.0], [0.0, 0.0]]])
x = {'f1': f1, 'f2': f2}
@tf.function
def model_fn(inp):
return tf.reduce_sum(input_tensor=w1 * inp['f1']) + tf.reduce_sum(
input_tensor=w2 * inp['f2'])
@tf.function
def loss_fn(_, pred):
return pred
adv_config = configs.AdvNeighborConfig(
feature_mask=None,
adv_step_size=0.1,
adv_grad_norm='l2',
iterations=2,
epsilon=0.15)
return x, w1, w2, adv_config, model_fn, loss_fn
def testVirtualAdvRegularizerMultiStepApproximation(self):
"""Tests virtual_adv_regularizer with multi-step approximation."""
np_input = np.array([[0.28, -0.96]])
tf_input = tf.constant(np_input)
embedding_fn = lambda x: x
vadv_config = configs.VirtualAdvConfig(
adv_neighbor_config=configs.AdvNeighborConfig(
feature_mask=None,
adv_step_size=1,
adv_grad_norm=configs.NormType.L2),
distance_config=configs.DistanceConfig(
distance_type=configs.DistanceType.COSINE, sum_over_axis=-1),
num_approx_steps=20,
approx_difference=1)
np_seed = np.array([[0.6, 0.8]])
tf_seed = tf.constant(np_seed)
vadv_loss = regularizer._virtual_adv_regularizer(
tf_input, embedding_fn, vadv_config, embedding_fn(tf_input),
tf_seed)
actual_loss = self.evaluate(vadv_loss)
# For detail derivation of the Hessian matrix, see go/vadv-tests-hessian
x = np_input
hessian = np.dot(x, x.T) * np.identity(2) - np.dot(x.T, x)
hessian /= np.linalg.norm(x)**4
approx = np.matmul(np_seed, hessian)
approx /= np.linalg.norm(approx, axis=-1, keepdims=True)
expected_loss = np.matmul(np.matmul(approx, hessian),
np.transpose(approx))
self.assertNear(actual_loss, expected_loss, err=1e-5)
def _test_gen_adv_neighbor_to_raise_for_disconnected_input_setup(self):
x = {
'f1': tf.constant([[1.0]]),
'f2': tf.constant([[2.0]]),
}
w = tf.constant([3.0])
adv_config = configs.AdvNeighborConfig(
feature_mask=None, adv_step_size=0.1, adv_grad_norm='l2')
return x, w, adv_config
def _test_gen_adv_neighbor_to_raise_for_nondifferentiable_input_setup(self):
x = {
'float': tf.constant([[-1.0]]),
'int': tf.constant([[2]], dtype=tf.dtypes.int32),
}
y = tf.constant([0.0])
w = tf.constant([[3.0], [4.0]])
adv_config = configs.AdvNeighborConfig(
feature_mask=None, adv_step_size=0.1, adv_grad_norm='l2')
return x, y, w, adv_config
def _test_gen_adv_neighbor_to_raise_for_sparse_tensors_setup(self):
sparse_feature = tf.SparseTensor(indices=[[0, 0], [0, 2], [1, 1]],
values=[1.0, 2.0, 3.0],
dense_shape=(2, 3))
x = {'sparse': sparse_feature, 'dense': tf.constant([[-1.0], [1.0]])}
w_sparse = tf.constant([[5.0], [4.0], [3.0]])
w_dense = tf.constant([[6.0]])
adv_config = configs.AdvNeighborConfig(feature_mask=None,
adv_step_size=0.1,
adv_grad_norm='l2')
return x, w_sparse, w_dense, adv_config
def _test_gen_adv_neighbor_for_features_with_different_shapes_setup(self):
w1 = tf.constant(1.0, shape=(2, 2, 3))
w2 = tf.constant(1.0, shape=(2, 2))
f1 = tf.constant([[[[0.0, 0.1, 0.2], [0.3, 0.4, 0.5]],
[[0.6, 0.7, 0.8], [0.9, 1.0, 1.1]]],
[[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]],
[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]]])
f2 = tf.constant([[[1.2, 1.3], [1.4, 1.5]], [[0.0, 0.0], [0.0, 0.0]]])
x = {'f1': f1, 'f2': f2}
adv_config = configs.AdvNeighborConfig(
feature_mask=None, adv_step_size=0.1, adv_grad_norm='l2')
return x, w1, w2, adv_config
def test_gen_adv_neighbor_respects_feature_constraints(
self, gen_adv_neighbor_fn):
x = tf.constant([[0.0, 1.0]])
w = tf.constant([[-1.0, 1.0]])
loss_fn = lambda x: tf.linalg.matmul(x, w, transpose_b=True)
adv_config = configs.AdvNeighborConfig(feature_mask=None,
adv_step_size=0.1,
adv_grad_norm='l2',
clip_value_min=0.0,
clip_value_max=1.0)
adv_neighbor = gen_adv_neighbor_fn(x, loss_fn, adv_config)
actual_neighbor = self.evaluate(adv_neighbor)
self.assertAllClose(x, actual_neighbor)
def _test_gen_adv_neighbor_for_feature_columns_setup(self):
# For linear regression
x = {
'fc1': tf.constant([[-1.0]]),
'fc2': tf.constant([[1.0]]),
}
y = tf.constant([0.0])
w = tf.constant([[3.0], [4.0]])
# gradient = [[6, 8]], normalized gradient = [[0.6, 0.8]]
expected_neighbor_fc1 = [[-1.0 + 0.1 * 0.6]]
expected_neighbor_fc2 = [[1.0 + 0.1 * 0.8]]
adv_config = configs.AdvNeighborConfig(
feature_mask={}, adv_step_size=0.1, adv_grad_norm='l2')
return x, y, w, expected_neighbor_fc1, expected_neighbor_fc2, adv_config
def test_gen_adv_neighbor_for_tensor_list(self):
x = [tf.constant([[-1.0]]), tf.constant([[1.0]])]
y = tf.constant([0.0])
w = [tf.constant([[3.0]]), tf.constant([[4.0]])]
with tf.GradientTape() as tape:
tape.watch(x)
y_hat = tf.matmul(x[0], w[0]) + tf.matmul(x[1], w[1])
loss = tf.math.squared_difference(y, y_hat)
adv_config = configs.AdvNeighborConfig(
feature_mask=tf.constant(1.0), adv_step_size=0.1, adv_grad_norm='l2')
adv_neighbor, _ = adv_lib.gen_adv_neighbor(
x, loss, adv_config, gradient_tape=tape)
# gradient = [[6, 8]], normalized gradient = [[0.6, 0.8]]
expected_neighbor = [[[-1.0 + 0.1 * 0.6]], [[1.0 + 0.1 * 0.8]]]
actual_neighbor = self.evaluate(adv_neighbor)
self.assertAllClose(expected_neighbor, actual_neighbor)
def test_multi_iter_gen_adv_neighbor_for_feature_columns_with_int_feature_v2(
self):
x, y, w, expected_neighbor_fc1, expected_neighbor_fc2 = (
self
._test_multi_iter_gen_adv_neighbor_for_feature_columns_with_int_feature(
))
loss_fn = tf.math.squared_difference
@tf.function
def model_fn(inp):
x_stacked = tf.concat(
[inp['fc1'], inp['fc2'],
tf.cast(inp['fc_int'], tf.dtypes.float32)],
axis=1)
return tf.matmul(x_stacked, w)
with tf.GradientTape() as tape:
tape.watch(x)
# Simple linear regression.
x_stacked = tf.concat(
[x['fc1'], x['fc2'],
tf.cast(x['fc_int'], tf.dtypes.float32)], axis=1)
y_hat = tf.matmul(x_stacked, w)
loss = tf.math.squared_difference(y, y_hat)
adv_config = configs.AdvNeighborConfig(
feature_mask=None,
adv_step_size=0.1,
adv_grad_norm='l2',
iterations=2,
epsilon=0.13)
adv_neighbor, _ = adv_lib.gen_adv_neighbor(
x,
loss,
adv_config,
gradient_tape=tape,
pgd_model_fn=model_fn,
pgd_loss_fn=loss_fn,
pgd_labels=y)
actual_neighbor = self.evaluate(adv_neighbor)
self.assertAllClose(expected_neighbor_fc1, actual_neighbor['fc1'])
self.assertAllClose(expected_neighbor_fc2, actual_neighbor['fc2'])
self.assertAllClose([[2]], actual_neighbor['fc_int'])
def test_gen_adv_neighbor_for_single_tensor_feature(self):
# Simple linear regression
x = tf.constant([[-1.0, 1.0]])
y = tf.constant([0.0])
w = tf.constant([[3.0], [4.0]])
with tf.GradientTape() as tape:
tape.watch(x)
y_hat = tf.matmul(x, w)
loss = tf.math.squared_difference(y, y_hat)
adv_config = configs.AdvNeighborConfig(
feature_mask=tf.constant(1.0), adv_step_size=0.1, adv_grad_norm='l2')
adv_neighbor, _ = adv_lib.gen_adv_neighbor(
x, loss, adv_config, gradient_tape=tape)
# gradient = [[6, 8]], normalized gradient = [[0.6, 0.8]]
expected_neighbor = [[-1.0 + 0.1 * 0.6, 1.0 + 0.1 * 0.8]]
actual_neighbor = self.evaluate(adv_neighbor)
self.assertAllClose(expected_neighbor, actual_neighbor)
def _test_multi_iter_gen_adv_neighbor_for_feature_columns_setup(self):
# For linear regression
x = {
'fc1': tf.constant([[-1.0]]),
'fc2': tf.constant([[1.0]]),
}
y = tf.constant([0.0])
w = tf.constant([[3.0], [4.0]])
# gradient = [[6, 8]], normalized gradient = [[0.6, 0.8]]
# Two iterations of 0.6 * 0.1, clipped.
expected_neighbor_fc1 = [[-1.0 + 0.13 * 0.6]]
# Two iterations of 0.8 * 0.1, clipped.
expected_neighbor_fc2 = [[1.0 + 0.13 * 0.8]]
adv_config = configs.AdvNeighborConfig(
feature_mask={},
adv_step_size=0.1,
adv_grad_norm='l2',
iterations=2,
epsilon=0.13)
return x, y, w, expected_neighbor_fc1, expected_neighbor_fc2, adv_config
def testVirtualAdvRegularizerRandomPerturbation(self):
"""Tests virtual_adv_regularizer with num_approx_steps=0."""
input_layer = tf.constant([[1.0, -1.0]])
embedding_fn = lambda x: x
step_size = 0.1
vadv_config = configs.VirtualAdvConfig(
adv_neighbor_config=configs.AdvNeighborConfig(
feature_mask=None,
adv_step_size=step_size,
adv_grad_norm=configs.NormType.L2),
distance_config=configs.DistanceConfig(
distance_type=configs.DistanceType.L2, sum_over_axis=-1),
num_approx_steps=0)
vadv_loss = regularizer.virtual_adv_regularizer(
input_layer, embedding_fn, vadv_config)
actual_loss = self.evaluate(vadv_loss)
# The identity embedding_fn makes the virtual adversarial loss immune to the
# direction of the perturbation, only the size matters.
expected_loss = step_size**2 # square loss
self.assertNear(actual_loss, expected_loss, err=1e-5)
def _gen_adv_neighbor(x):
w = tf.constant([[3.0], [4.0]])
loss = tf.matmul(x, w)
y = tf.constant([0.0])
model_fn = lambda inp: tf.matmul(inp, w)
@tf.function
def loss_fn(_, pred):
return pred
adv_config = configs.AdvNeighborConfig(
feature_mask=None,
adv_step_size=0.1,
adv_grad_norm='l2',
epsilon=0.15,
iterations=2)
adv_neighbor, _ = adv_lib.gen_adv_neighbor(
x,
loss,
adv_config,
pgd_labels=y,
pgd_model_fn=model_fn,
pgd_loss_fn=loss_fn)
return adv_neighbor
def _test_gen_adv_neighbor_for_all_input_disconnected_setup(self):
x = tf.constant([[1.0]])
loss = tf.constant(1.0)
adv_config = configs.AdvNeighborConfig()
return x, loss, adv_config