Implementation of Adversarial Loss In Keras

Question

I'm trying to implement an adversarial loss in keras. The model consists of two networks, one auto-encoder (the target model) and one discriminator. The two models share the encoder.

I created the adversarial loss of the auto-encoder by setting a keras variable

def get_adv_loss(d_loss):
    def loss(y_true, y_pred):
        return some_loss(y_true, y_pred) - d_loss
    return loss

discriminator_loss = K.variable()
L = get_adv_loss(discriminator_loss)
autoencoder.compile(..., loss=L)

and during training I interleave train_on_batch of discriminator and autoencoder to update discriminator_loss

d_loss = disciminator.train_on_batch(x, y_domain)
discriminator_loss.assign(d_loss)
a_loss, ... = self.segmenter.train_on_batch(x, y_target)

However, I found out that the value of these variables is frozen when the model is compiled. I tried to recompile the model during training but that raise the error

Node 'IsVariableInitialized_13644': Unknown input node 'training_12/Adam/Variable'

which I guess it means i cant recompile during training? any suggestion on how i can inject the discriminator loss in the autoencoder?

Answer 1

Keras model supports multiple outputs. So just include your discirminator into your keras model and freeze the discrminator layers, if the discriminator should not be trained.

The next question would be how to combine autoencoder loss and discriminator loss. Luckily keras model.compile supports loss weights. If autoencoder is your first output and discriminator is your second you could do something like loss_weights=[1, -1]. So a better discriminator is worse for the autoencoder.

Edit: Here is an example, how to implement an Adversary Network:

# Build your architecture
auto_encoder_input = Input((5,))
auto_encoder_net = Dense(10)(auto_encoder_input)
auto_encoder_output = Dense(5)(auto_encoder_net)

discriminator_net = Dense(20)(auto_encoder_output)
discriminator_output = Dense(5)(discriminator_net)

# Define outputs of your model
train_autoencoder_model = Model(auto_encoder_input, [auto_encoder_output, discriminator_output])
train_discriminator_model = Model(auto_encoder_input, discriminator_output)

# Compile the models (compile the first model and then change the trainable attribute for the second)
for layer_index, layer in enumerate(train_autoencoder_model.layers):
    layer.trainable = layer_index < 3

train_autoencoder_model.compile('Adam', loss=['mse', 'mse'], loss_weights=[1, -1])        

for layer_index, layer in enumerate(train_discriminator_model.layers):
    layer.trainable = layer_index >= 3

train_discriminator_model.compile('Adam', loss='mse')

# A simple example how a training can look like
for i in range(10):
    auto_input = np.random.sample((10,5))
    discrimi_output = np.random.sample((10,5))
    train_discriminator_model.fit(auto_input, discrimi_output, steps_per_epoch=5, epochs=1)
    train_autoencoder_model.fit(auto_input, [auto_input, discrimi_output], steps_per_epoch=1, epochs=1)  

As you can see there is no much magic behind building an Adversary Model with keras.