How to apply regularization when fine tuning on the imagenet extracted models

Viewed 720

One way of fine tuning is extracting a model ( like VGG16 trained on Imagenet) , adding a layer or so and then train the model.

Is it possible to apply regularization to the model layers apart from the added layer using Tensorflow.Keras. I don't think adding regularization to only one layer effects the outcome much.

I know we can apply the regularization for the added layer as:

x = Dense(classes, kernel_regularizer=l2(reg), name="labels")(x)

But is it possible to apply regularization for other layers as well in Keras. It could be easily done in mxnet.

Would be grateful for any help.

3 Answers

This solution should work. By iterating over the model layers we can just add a regularizer. You can then add your dense layer after this.

model = tf.keras.applications.VGG16(include_top=False, weights=None)
regularizer = tf.keras.regularizers.l2(0.001)

for i in range(len(model.layers)):
    if isinstance(model.layers[i], tf.keras.layers.Conv2D):
        print('Adding regularizer to layer {}'.format(model.layers[i].name))
        model.layers[i].kernel_regularizer = regularizer

# Add Dense layer
classes = 10
x = model.output
x = Dense(classes, kernel_regularizer=regularizer, name="labels")(x)

model = tf.keras.Model(model.input, x)

Unfortunately @DMolony's answer doesn't update the graph, meaning that your model will not apply regularization to the changed layers.

You can examine model.losses to check if the loss has been properly updated. Running @DMolony's snippet:

import tensorflow as tf

model = tf.keras.applications.VGG16(include_top=False, weights=None)
regularizer = tf.keras.regularizers.l2(0.001)
print(f'model.losses (original): {model.losses}')

for i in range(len(model.layers)):
  if isinstance(model.layers[i], tf.keras.layers.Conv2D):
    print('Adding regularizer to layer {}'.format(model.layers[i].name))
    model.layers[i].kernel_regularizer = regularizer
print(f'model.losses (after setting reg): {model.losses}')

# Add Dense layer
classes = 10
x = model.output
x = tf.keras.layers.Dense(
    classes, kernel_regularizer=regularizer, name='labels')(
        x)

model = tf.keras.Model(model.input, x)
print(f'model.losses (after adding dense): {model.losses}')

We see that the new regularization losses are not added to the model:

model.losses (original): []
Adding regularizer to layer block1_conv1
Adding regularizer to layer block1_conv2
...
Adding regularizer to layer block5_conv3
model.losses (after setting reg): []
model.losses (after adding dense): [<tf.Tensor: shape=(), dtype=float32, numpy=0.019377343>]

If you want to update the graph, you needed to 1) save the updated model config, 2) save the current weights, and 3) reload the config and weights after adding regularizers:

import os
import tempfile
import tensorflow as tf

model = tf.keras.applications.VGG16(include_top=False, weights=None)
regularizer = tf.keras.regularizers.l2(0.001)
print(f'model.losses (original): {model.losses}')

for layer in model.layers:
  if isinstance(layer, tf.keras.layers.Conv2D) and hasattr(layer, 'kernel_regularizer'):
    print('\tAdding regularizer to layer {}'.format(layer.name))
    setattr(layer, 'kernel_regularizer', regularizer)
print(f'model.losses (after setting reg): {model.losses}')

# Add Dense layer
classes = 10
x = model.output
x = tf.keras.layers.Dense(
    classes, kernel_regularizer=regularizer, name='labels')(
        x)

model = tf.keras.Model(model.input, x)
print(f'model.losses (after adding dense): {model.losses}')

tmp_weights_dir = tempfile.gettempdir()
tmp_weights_path = os.path.join(tmp_weights_dir, 'tmp_weights.h5')

# Save model config and weights and reload the model from these values.
model_json = model.to_json()
model.save_weights(tmp_weights_path)
model = tf.keras.models.model_from_json(model_json)
model.load_weights(tmp_weights_path, by_name=True)
print(f'model.losses (after reloading):')
loss_str = '\n'.join(['\t' + str(loss) for loss in model.losses])
print(loss_str)

Now the losses are available in the model:

odel.losses (original): []
    Adding regularizer to layer block1_conv1
    Adding regularizer to layer block1_conv2
    Adding regularizer to layer block2_conv1
    Adding regularizer to layer block2_conv2
    Adding regularizer to layer block3_conv1
    Adding regularizer to layer block3_conv2
    Adding regularizer to layer block3_conv3
    Adding regularizer to layer block4_conv1
    Adding regularizer to layer block4_conv2
    Adding regularizer to layer block4_conv3
    Adding regularizer to layer block5_conv1
    Adding regularizer to layer block5_conv2
    Adding regularizer to layer block5_conv3
model.losses (after setting reg): []
model.losses (after adding dense): [<tf.Tensor: shape=(), dtype=float32, numpy=0.019144716>]
model.losses (after reloading):
    tf.Tensor(0.0058215745, shape=(), dtype=float32)
    tf.Tensor(0.06401116, shape=(), dtype=float32)
    tf.Tensor(0.08538648, shape=(), dtype=float32)
    tf.Tensor(0.12775433, shape=(), dtype=float32)
    tf.Tensor(0.17033958, shape=(), dtype=float32)
    tf.Tensor(0.2565217, shape=(), dtype=float32)
    tf.Tensor(0.25641307, shape=(), dtype=float32)
    tf.Tensor(0.34119228, shape=(), dtype=float32)
    tf.Tensor(0.5120113, shape=(), dtype=float32)
    tf.Tensor(0.51162773, shape=(), dtype=float32)
    tf.Tensor(0.5116309, shape=(), dtype=float32)
    tf.Tensor(0.5124292, shape=(), dtype=float32)
    tf.Tensor(0.5116511, shape=(), dtype=float32)
    tf.Tensor(0.019144716, shape=(), dtype=float32)

There is actually a solution without reloading the model graph and its weights!

You can leverage the public keras api with the Layer.add_loss instance method, which will automatically keeps track of dependencies for you.

I updated the example with this api:

import tensorflow as tf

model = tf.keras.applications.VGG16(include_top=False, weights=None)
regularizer = tf.keras.regularizers.l2(1e-3)
print('\nmodel.losses (original): \n', model.losses)

for lyr in model.layers:
  if isinstance(lyr, tf.keras.layers.Conv2D):
    print(f'Adding regularizer to layer {lyr.name}')
    lyr.add_loss(lambda x=lyr.kernel: regularizer(x))
print('\nmodel.losses (after setting reg): \n', model.losses)

# Add Dense layer
classes = 10
x = model.output
x = tf.keras.layers.Dense(
    classes, kernel_regularizer=regularizer, name='labels')(x)

model = tf.keras.Model(model.input, x)
print('\nmodel.losses (after adding dense): \n', model.losses)
Related