def test_variable(self): """Test invoking Variable.""" value = np.random.rand(5, 4).astype(np.float32) layer = layers.Variable(value) layer.build([]) result = layer.call([]).numpy() assert np.allclose(result, value) assert len(layer.trainable_variables) == 1
def add_adapter(self, all_layers, task, layer_num):
"""Add an adapter connection for given task/layer combo"""
i = layer_num
prev_layers = []
trainable_layers = []
# Handle output layer
if i < len(self.layer_sizes):
layer_sizes = self.layer_sizes
alpha_init_stddev = self.alpha_init_stddevs[i]
weight_init_stddev = self.weight_init_stddevs[i]
bias_init_const = self.bias_init_consts[i]
elif i == len(self.layer_sizes):
layer_sizes = self.layer_sizes + [self.n_outputs]
alpha_init_stddev = self.alpha_init_stddevs[-1]
weight_init_stddev = self.weight_init_stddevs[-1]
bias_init_const = self.bias_init_consts[-1]
else:
raise ValueError("layer_num too large for add_adapter.")
# Iterate over all previous tasks.
for prev_task in range(task):
prev_layers.append(all_layers[(i - 1, prev_task)])
# prev_layers is a list with elements of size
# (batch_size, layer_sizes[i-1])
if len(prev_layers) == 1:
prev_layer = prev_layers[0]
else:
prev_layer = Concatenate(axis=1)(prev_layers)
alpha = layers.Variable(
tf.random.truncated_normal((1, ), stddev=alpha_init_stddev))
trainable_layers.append(alpha)
prev_layer = Multiply()([prev_layer, alpha([prev_layer])])
dense1 = Dense(
layer_sizes[i - 1],
kernel_initializer=tf.keras.initializers.TruncatedNormal(
stddev=weight_init_stddev),
bias_initializer=tf.constant_initializer(value=bias_init_const))
prev_layer = dense1(prev_layer)
trainable_layers.append(dense1)
dense2 = Dense(
layer_sizes[i],
kernel_initializer=tf.keras.initializers.TruncatedNormal(
stddev=weight_init_stddev),
use_bias=False)
prev_layer = dense2(prev_layer)
trainable_layers.append(dense2)
return prev_layer, trainable_layers
def __init__(self, n_generators=1, n_discriminators=1, **kwargs):
"""Construct a GAN.
In addition to the parameters listed below, this class accepts all the
keyword arguments from KerasModel.
Parameters
----------
n_generators: int
the number of generators to include
n_discriminators: int
the number of discriminators to include
"""
self.n_generators = n_generators
self.n_discriminators = n_discriminators
# Create the inputs.
self.noise_input = Input(shape=self.get_noise_input_shape())
self.data_input_layers = []
for shape in self.get_data_input_shapes():
self.data_input_layers.append(Input(shape=shape))
self.data_inputs = [i.ref() for i in self.data_input_layers]
self.conditional_input_layers = []
for shape in self.get_conditional_input_shapes():
self.conditional_input_layers.append(Input(shape=shape))
self.conditional_inputs = [
i.ref() for i in self.conditional_input_layers
]
# Create the generators.
self.generators = []
self.gen_variables = []
generator_outputs = []
for i in range(n_generators):
generator = self.create_generator()
self.generators.append(generator)
generator_outputs.append(
generator(
_list_or_tensor([self.noise_input] +
self.conditional_input_layers)))
self.gen_variables += generator.trainable_variables
# Create the discriminators.
self.discriminators = []
self.discrim_variables = []
discrim_train_outputs = []
discrim_gen_outputs = []
for i in range(n_discriminators):
discriminator = self.create_discriminator()
self.discriminators.append(discriminator)
discrim_train_outputs.append(
self._call_discriminator(discriminator, self.data_input_layers,
True))
for gen_output in generator_outputs:
if isinstance(gen_output, tf.Tensor):
gen_output = [gen_output]
discrim_gen_outputs.append(
self._call_discriminator(discriminator, gen_output, False))
self.discrim_variables += discriminator.trainable_variables
# Compute the loss functions.
gen_losses = [
self.create_generator_loss(d) for d in discrim_gen_outputs
]
discrim_losses = []
for i in range(n_discriminators):
for j in range(n_generators):
discrim_losses.append(
self.create_discriminator_loss(
discrim_train_outputs[i],
discrim_gen_outputs[i * n_generators + j]))
if n_generators == 1 and n_discriminators == 1:
total_gen_loss = gen_losses[0]
total_discrim_loss = discrim_losses[0]
else:
# Create learnable weights for the generators and discriminators.
gen_alpha = layers.Variable(np.ones((1, n_generators)),
dtype=tf.float32)
# We pass an input to the Variable layer to work around a bug in TF 1.14.
gen_weights = Softmax()(gen_alpha([self.noise_input]))
discrim_alpha = layers.Variable(np.ones((1, n_discriminators)),
dtype=tf.float32)
discrim_weights = Softmax()(discrim_alpha([self.noise_input]))
# Compute the weighted errors
weight_products = Reshape(
(n_generators * n_discriminators, ))(Multiply()([
Reshape((n_discriminators, 1))(discrim_weights),
Reshape((1, n_generators))(gen_weights)
]))
stacked_gen_loss = layers.Stack(axis=0)(gen_losses)
stacked_discrim_loss = layers.Stack(axis=0)(discrim_losses)
total_gen_loss = Lambda(lambda x: tf.reduce_sum(x[0] * x[1]))(
[stacked_gen_loss, weight_products])
total_discrim_loss = Lambda(lambda x: tf.reduce_sum(x[0] * x[1]))(
[stacked_discrim_loss, weight_products])
self.gen_variables += gen_alpha.trainable_variables
self.discrim_variables += gen_alpha.trainable_variables
self.discrim_variables += discrim_alpha.trainable_variables
# Add an entropy term to the loss.
entropy = Lambda(lambda x: -(tf.reduce_sum(tf.math.log(x[
0])) / n_generators + tf.reduce_sum(tf.math.log(x[
1])) / n_discriminators))([gen_weights, discrim_weights])
total_discrim_loss = Lambda(lambda x: x[0] + x[1])(
[total_discrim_loss, entropy])
# Create the Keras model.
inputs = [self.noise_input
] + self.data_input_layers + self.conditional_input_layers
outputs = [total_gen_loss, total_discrim_loss]
self.gen_loss_fn = lambda outputs, labels, weights: outputs[0]
self.discrim_loss_fn = lambda outputs, labels, weights: outputs[1]
model = tf.keras.Model(inputs=inputs, outputs=outputs)
super(GAN, self).__init__(model, self.gen_loss_fn, **kwargs)