def test_gradient_penalty_loss_with_wrong_input_types_raises(self): discriminator = tf.keras.Sequential() with self.assertRaisesRegex( TypeError, 'should either both be a tf.Tensor ' 'or both a sequence of tf.Tensor'): losses.gradient_penalty_loss(real_data=(tf.ones((1, )), ), generated_data=tf.ones((1, )), discriminator=discriminator)
def test_gradient_penalty_loss_with_unequal_number_of_elements_raises( self): discriminator = tf.keras.Sequential() with self.assertRaisesRegex( ValueError, 'number of elements in real_data and generated_data are ' 'expected to be equal'): losses.gradient_penalty_loss(real_data=(tf.ones((1, )), ), generated_data=(tf.ones( (1, )), tf.ones((1, ))), discriminator=discriminator)
def test_gradient_penalty_loss_positive(self): discriminator = tf.keras.Sequential() discriminator.add(tf.keras.layers.Reshape((25, ))) discriminator.add(tf.keras.layers.Dense(units=1)) real_data = tf.ones(shape=(1, 5, 5)) generated_data = tf.ones(shape=(1, 5, 5)) gradient_penalty = losses.gradient_penalty_loss( real_data=real_data, generated_data=generated_data, discriminator=discriminator) self.assertAllGreaterEqual(gradient_penalty, 0.0)
def test_gradient_penalty_shape_correct(self): discriminator = tf.keras.Sequential() discriminator.add(tf.keras.layers.Reshape((25, ))) discriminator.add(tf.keras.layers.Dense(units=1)) real_data = tf.ones(shape=(3, 5, 5)) generated_data = tf.ones(shape=(3, 5, 5)) gradient_penalty = losses.gradient_penalty_loss( real_data=real_data, generated_data=generated_data, discriminator=discriminator) self.assertAllEqual(tf.shape(gradient_penalty), (3, ))
def gradient_penalty_fn(weights): def multiply(input_tensor): return tf.linalg.matmul(input_tensor, weights) discriminator = tf.keras.Sequential() discriminator.add(tf.keras.layers.Reshape((25, ))) # To simulate a dense layer a lambda layer is used, such that we are able # to feed the weights in as numpy array to the assert_jacobian_fn. discriminator.add(tf.keras.layers.Lambda(multiply)) return losses.gradient_penalty_loss( real_data=tf.convert_to_tensor(real_data), generated_data=tf.convert_to_tensor(generated_data), discriminator=discriminator)
def test_gradient_penalty_loss_lambda_for_zero_gradient(self): discriminator = tf.keras.Sequential() discriminator.add(tf.keras.layers.Reshape((4, ))) # Generates a dense layer that is initialized with all zeros. # This leads to a network that has zero gradient for any input. discriminator.add( tf.keras.layers.Dense(units=1, kernel_initializer='zeros', bias_initializer='zeros')) real_data = tf.ones(shape=(1, 2, 2)) generated_data = tf.ones(shape=(1, 2, 2)) weight = 1.0 gradient_penalty = losses.gradient_penalty_loss( real_data=real_data, generated_data=generated_data, discriminator=discriminator, weight=weight) # Tolerance is large due to eps that is added in the gradient pentaly loss # for numerical stability at 0. self.assertAllClose(gradient_penalty, (weight, ), atol=0.001)