Find input that maximises output of a neural network using Keras and TensorFlow

Question

I have used Keras and TensorFlow to classify the Fashion MNIST following this tutorial .

It uses the AdamOptimizer to find the value for model parameters that minimize the loss function of the network. The input for the network is a 2-D tensor with shape [28, 28], and output is a 1-D tensor with shape [10] which is the result of a softmax function.

Once the network has been trained, I want to use the optimizer for another task: find an input that maximizes one of the elements of the output tensor. How can this be done? Is it possible to do so using Keras or one have to use a lower level API?

Since the input is not unique for a given output, it would be even better if we could impose some constraints on the values the input can take.

The trained model has the following format

model = keras.Sequential([
    keras.layers.Flatten(input_shape=(28, 28)),
    keras.layers.Dense(128, activation=tf.nn.relu),
    keras.layers.Dense(10, activation=tf.nn.softmax)
])

Answer 1

I feel you would want to backprop with respect to the input freezing all the weights to your model. What you could do is:

1. Add a dense layer after the input layer with the same dimensions as input and set it as trainable
2. Freeze all the other layers of your model. (except the one you added)
3. As an input, feed an identity matrix and train your model based on whatever output you desire.

This article and this post might be able to help you if you want to backprop based on the input instead. It's a bit like what you are aiming for but you can get the intuition.

Answer 2

It would be very similar to the way that filters of a Convolutional Network is visualized: we would do gradient ascent optimization in input space to maximize the response of a particular filter.

Here is how to do it: after training is finished, first we need to specify the output and define a loss function that we want to maximize:

from keras import backend as K

output_class = 0 # the index of the output class we want to maximize
output = model.layers[-1].output
loss = K.mean(output[:,output_class]) # get the average activation of our desired class over the batch

Next, we need to take the gradient of the loss we have defined above with respect to the input layer:

grads = K.gradients(loss, model.input)[0] # the output of `gradients` is a list, just take the first (and only) element

grads = K.l2_normalize(grads) # normalize the gradients to help having an smooth optimization process

Next, we need to define a backend function that takes the initial input image and gives the values of loss and gradients as outputs, so that we can use it in the next step to implement the optimization process:

func = K.function([model.input], [loss, grads])

Finally, we implement the gradient ascent optimization process:

import numpy as np

input_img = np.random.random((1, 28, 28)) # define an initial random image

lr = 1.  # learning rate used for gradient updates
max_iter = 50  # number of gradient updates iterations
for i in range(max_iter):
    loss_val, grads_val = func([input_img])
    input_img += grads_val * lr  # update the image based on gradients

Note that, after this process is finished, to display the image you may need to make sure that all the values in the image are in the range [0, 255] (or [0,1]).