Unable to completely separate outputs of model in TensorFlow

user_2514130 0
Viewed 124

I am trying to create a convolutional neural network that has two regression outputs, a score and a confidence. I have frozen the layers they have in common in the hopes that the addition of the confidence output doesn't change the score, but in my experiments it has. For the model with just the score, I used Xception and added a simple GlobalAveragePooling2D and Dense(512) layer then output a single number.

base_model = Xception(input_shape=(224, 224, 3), weights='imagenet', include_top=False)

x = base_model.output
x = GlobalAveragePooling2D()(x)
x = Dense(512, activation='relu')(x)
predictions = Dense(1, activation='sigmoid')(x)

model = Model(inputs=base_model.input, outputs=predictions)

for layer in base_model.layers:
    layer.trainable = False

optimizer = Adam(learning_rate=learning_rate)
model.compile(loss='mae', optimizer=optimizer, metrics=['mse','mae'], run_eagerly=True)

Here is what the end of model.summary() looks like:

enter image description here

When I fit it, the model produces good results.

But when I try to add a second output the result of the first becomes much worse. The new model gets trained off tuples where is first number is the same as the first model and the second number is a confidence value. The model is very similar to the one above.

base_model = Xception(input_shape=(224, 224, 3), weights='imagenet', include_top=False)

x = base_model.output
x = GlobalAveragePooling2D()(x)

score_x = Dense(512, activation='relu')(x)
score_out = Dense(1, activation='sigmoid', name='score_model')(score_x)

confidence_x = Dense(512, activation='relu')(x)
confidence_out = Dense(1, name='confidence_model')(confidence_x)

model = Model(inputs=base_model.input, outputs=[score_out, confidence_out])

for layer in base_model.layers:
    layer.trainable = False

losses = {'score_model': 'mae', 'confidence_model': 'mae'}
loss_weights = {'score_model': 1, 'confidence_model': 1} 

model.compile(loss=losses, loss_weights=loss_weights, optimizer=optimizer, metrics=['mse','mae'], run_eagerly=True)

When I look at model.summary(), it has twice as many trainable parameters as the previous model, which is exactly what I was expecting. Everything looks right to me so far.

enter image description here

But when I train this model the performance on the score is much worse. I was thinking it would be the same (within stochastic variation). After the first epoch, the loss from the first model is around 0.125. The score_model_loss from the second model is around 0.554. Clearly I'm not completely separating the models. What am I missing?

conv-neural-network keras tensorflow python
1 Answers

Note: This answer will work well only because the layer that do the feature extraction are frozen. As @Akshay Sehgal stated in the comments :

optimizing for 2 goals together is actually a completely different problem than optimizing 2 independent goals separately

In that case, we are optimizing for 2 goals separately.

The easiest solution is probably to write a custom training loop with 2 tf.GradientTape, one for each goal. Lets consider this really simple example:

Dummy data

Let's create some random Data

import tensorflow as tf

X = tf.random.normal((1000,1))
y1= 3*X + 1
y2 = -2*X +2

ds = tf.data.Dataset.from_tensor_slices((X,y1,y2)).batch(10)

Creating a model with 2 outputs

In that example, I skip the feature extraction step, as a simple linear regression will work for the data. But as your feature extractor network is frozen, the example is similar.

inp = tf.keras.Input((1,))
dense_1 = tf.keras.layers.Dense(1, name="objective1")(inp)
dense_2 = tf.keras.layers.Dense(1, name="objective2")(inp)
model = tf.keras.Model(inputs=inp, outputs=[dense_1, dense_2])

# setting up the loss functions as well as the optimizer
opt = tf.optimizers.SGD()
loss_func1 = tf.losses.mean_squared_error
loss_func2 = tf.losses.mean_absolute_error

Note the name given to the two dense layers: I will use them later to retrieve the appropriate weights.

Getting the weights to optimize

We can use the name set before to retrieve the variable belonging to each objective :

var1, var2 = [],[]
for l in model.layers:
    if "objective1" in l.name:
        var1 += l.trainable_variables
    if "objective2" in l.name:
        var2 += l.trainable_variables

The training loop

You simply need to tapes, one for each objective. You can use different optimizer as well, if it makes the training better.

counter = 0
for x, y1, y2 in ds:
    counter += 1
    with tf.GradientTape() as tape1, tf.GradientTape() as tape2:
        pred1, pred2 = model(x)
        loss1 = loss_func1(y1, pred1)
        loss2 = loss_func2(y2, pred2)
    grad1 = tape1.gradient(loss1, var1)
    grad2 = tape2.gradient(loss2, var2)
    opt.apply_gradients(zip(grad1, var1))
    opt.apply_gradients(zip(grad2, var2))
    if counter % 10:
        print(f"Step : {counter}, objective1: {tf.reduce_mean(loss1)}, objective2: {tf.reduce_mean(loss2)}")

If we run the training, we get:

Step : 1, objective1: 4.609124183654785, objective2: 2.6634981632232666
[...]
Step : 99, objective1: 7.176481902227555e-14, objective2: 0.030187154188752174

The principle advantage training that way is that you just need to extract the features once for the two objectives.

user_73701530
Related