Transfer Learning/Fine Tuning - how to keep BatchNormalization in inference mode?

user_11687201 0
Viewed 406

In the following tutorial Transfer learning and fine-tuning by TensorFlow it is explained that that when unfreezing a model that contains BatchNormalization (BN) layers, these should be kept in inference mode by passing training=False when calling the base model.

[…]

Important notes about BatchNormalization layer

Many image models contain BatchNormalization layers. That layer is a special case on every imaginable count. Here are a few things to keep in mind.

  • BatchNormalization contains 2 non-trainable weights that get updated during training. These are the variables tracking the mean and variance of the inputs.
  • When you set bn_layer.trainable = False, the BatchNormalization layer will run in inference mode, and will not update its mean & variance statistics. This is not the case for other layers in general, as weight trainability & inference/training modes are two orthogonal concepts. But the two are tied in the case of the BatchNormalization layer.
  • When you unfreeze a model that contains BatchNormalization layers in order to do fine-tuning, you should keep the BatchNormalization layers in inference mode by passing training=False when calling the base model. Otherwise the updates applied to the non-trainable weights will suddenly destroy what the model has learned.

[…]

In the examples they pass training=False when calling the base model, but later they set base_model.trainable=True, which for my understanding is the opposite of inference mode, because the BN layers will be set to trainable as well.

For my understanding there would have to be 0 trainable_weights and 4 non_trainable_weights for inference mode, which would be identical to when setting the bn_layer.trainable=False, which they stated would be the case for running the bn_layer in inference mode.

I checked the number of trainable_weights and number of non_trainable_weights and they are both 2.

I am confused by the tutorial, how can I really be sure BN layer are in inference mode when doing fine tuning on a model?

Does setting training=False on the model overwrite the behavior of bn_layer.trainable=True? So that even if the trainable_weights get listed with 2 these would not get updated during training (fine tuning)?

Update:

Here I found some further information: BatchNormalization layer - on keras.io.

[...]

About setting layer.trainable = False on a BatchNormalization layer:

The meaning of setting layer.trainable = False is to freeze the layer, i.e. its internal state will not change during training: its trainable weights will not be updated during fit() or train_on_batch(), and its state updates will not be run.

Usually, this does not necessarily mean that the layer is run in inference mode (which is normally controlled by the training argument that can be passed when calling a layer). "Frozen state" and "inference mode" are two separate concepts.

However, in the case of the BatchNormalization layer, setting trainable = False on the layer means that the layer will be subsequently run in inference mode (meaning that it will use the moving mean and the moving variance to normalize the current batch, rather than using the mean and variance of the current batch).

This behavior has been introduced in TensorFlow 2.0, in order to enable layer.trainable = False to produce the most commonly expected behavior in the convnet fine-tuning use case.

Note that: - Setting trainable on an model containing other layers will recursively set the trainable value of all inner layers. - If the value of the trainable attribute is changed after calling compile() on a model, the new value doesn't take effect for this model until compile() is called again.

Question:

  1. In case I want to fine tune the whole model, so I am going to unfreeze the base_model.trainable = True, would I have to manually set the BN layers to bn_layer.trainable = False in order to keep them in inference mode?
  2. What does happen when with the call of the base_model passing training=False and additionally setting base_model.trainable=True? Do layers like BatchNormalization and Dropout stay in inference mode?
tensorflow2.0 tf.keras keras tensorflow
2 Answers

After reading the documentation and having a look on the source code of TensorFlows implementations of tf.keras.layers.Layer, tf.keras.layers.Dense, and tf.keras.layers.BatchNormalization I got the following understanding.

If training = False is passed on calling the layer or the model/base model, it will run in inference mode. This has nothing to do with the attribute trainable, which means something different. It would probably lead to less misunderstanding, if they would have called the parameter training_mode instead of training. I would have preferred defining it the other way round and calling it inference_mode .

When doing Transfer Learning or Fine Tuning training = False should be passed on calling the base model itself. As far as I saw until now this will only affect layers like tf.keras.layers.Dropout and tf.keras.layers.BatchNormalization and will have not effect on the other layers. Running in inference mode via training = False will result in:

  • tf.layers.Dropout not to apply the dropout rate at all. As tf.layers.Dropout has no trainable weights, setting the attribute trainable = False will have no effect at all on this layer.
  • tf.keras.layers.BatchNormalization normalizing its inputs using the mean and variance of its moving statistics learned during training

The attribute trainable will only activate or deactivate updating the trainable weights of a layer.

user_116872010

Yes, After calling base_model.trainable = True, you need to set the BN layers to inference mode.

from tensorflow.keras import layers
for layer in base_model.layers:
    if isinstance(layer, layers.BatchNormalization):
        layer.trainable = False

And you can check each layer, whether its in training or inference mode:

for lnum, layer in enumerate(model.layers):
    print('layer: {}, name:{}, trainable:{}, dtype: {}'.format(lnum, layer.name, layer.trainable, layer.dtype))
user_109383520
Related