I would like to implement an adversarial network with a classifier whose output is connected to an adversary that has to guess a specific feature of the inputs to the classifier (a nuisance parameter) based on the output of the classifier (a detailed description of such an adversarial network can be found in this paper: learning to pivot with adversarial networks
The model will then be trained like this:
- train the adversary on a batch (the classifier will be fixed)
- freeze the adversary
- train the whole model while the adversary is frozen on a batch
- unfreeze the adversary
- repeat for the other batches
- shuffle and repeat for other epochs
When training the adversary I want the loss function to be categorical crossentropy, and when training the whole model, I want it to be the loss function of the classifier (a binary crossentropy loss) minus the loss of the adversary times a parameter: L_c() - b * L_a().
Most of the adversarial training codes that I've seen are trained in Keras using train_on_batch. However since I already have a lot of code setup with another sequential model, and I would like to reuse it, I was wondering if there is a way to implement this model and train it with model.fit in Keras.
What I was thinking of doing was to setup the model using Keras functional API with the classifier input as its input and the classifier and adversary outputs as its outputs. I would also compile the model as an adversary model with only a single output. For example:
classifier_input = Input(shape=(10,))
x = Dense(50, activation='tanh')(classifier_input)
x = Dense(25, activation='tanh')(x)
x = Dense(5, activation='tanh')(x)
classifier_output = Dense(1, activation='sigmoid')(x)
x = Dense(30, activation='tanh')(classifier_output)
x = Dense(15, activation='tanh')(x)
x = Dense(5, activation='tanh')(x)
adversary_output = Dense(3, activation='sgimoid')(x)
adversary = Model(inputs=classifier_input , outputs=adversary_output)
adversary.compile(optimizer='Adamax', loss='categorical_crossentropy', metrics=['accuracy'])
final_model = Model(inputs=classifier_input,outputs=[classifier_output,adversary_output])
final_model.compile(optimizer='Adamax',
loss=['binary_crossentropy', 'categorical_crossentropy'],
loss_weights=[1.0, -0.1], metrics=['accuracy'])
And then I want to setup a Callback that will train the adversary (after freezing the layers in the classifier) in on_batch_begin, and then the training of final_model will be done using model.fit code (I will freeze adversary and unfreeze the classifier layers in on_batch_begin right before the training of final_model begins).
But then I don't know if it is possible to pass the current batch as an argument to on_batch_begin. Would I have to setup my own batches in the Callback or can the batch be passed by model.fit?
Is there a better way of doing adversarial training while still using model.fit?
