How to update weights in Stochastic Weight Averaging (SWA) on tensorflow?

user_607528 0
Viewed 927

I'm confused about how to implement tfa's SWA optimizer. There are two points here:

  1. When you look at the documentation it points you to [this] model averaging tutorial. That tutorial uses tfa.callbacks.AverageModelCheckpoint, which allows you to
  • Assign the moving average weights to the model, and save them.
  • (or) Keep the old non-averaged weights, but the saved model uses the average weights.

Having a distinct ModelCheckpoint that allows you to save moving average weights (rather than the current weights) makes sense. However - it seems like SWA should be managing the weight averaging. That makes me want to set update_weights=False.

Is this correct? The tutorial uses update_weights=True.

  1. There is a note about SWA not updating the BN layers in the documentation. Following the suggestion here I did this,
# original training
model.fit(...)
# updating weights from final run 
optimizer.assign_average_vars(model.variables)
# batch-norm-hack: lr=0 as suggested https://stackoverflow.com/a/64376062/607528
model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=0),
    loss=loss,
    metrics=metrics)
model.fit(
    data,
    validation_data=None,
    epochs=1,
    callbacks=final_callbacks)

before saving my model.

Is this correct?

Thanks!

tensorflow2.0 tensorflow deep-learning machine-learning python
2 Answers

The easiest way to deal with the batch norm is the following: First, loop through all layers in your model and reset the moving mean and moving variance in the batch norm layers (in my example I assume the batch norm layers end with "bn"):

for l in model.layers:
    if l.name.split('_')[-1] == 'bn': # e.g. conv1_bn
       l.moving_mean.assign(tf.zeros_like(l.moving_mean))
       l.moving_variance.assign(tf.ones_like(l.moving_variance))

After that run your model for one epoch and set training to true to update the moving average and variance:

count = 0
for x,_ in dataset_train:
    _ = model(x, training = True)
    count += 1
    if count > steps_per_epoch:
       break
user_154824630

There are two ways of doing this, the first one is you manually update the weights before saving, like this example from the documentation.

import tensorflow as tf
import tensorflow_addons as tfa

model = tf.Sequential([...])
opt = tfa.optimizers.SWA(
        tf.keras.optimizers.SGD(lr=2.0), 100, 10)
model.compile(opt, ...)
model.fit(x, y, ...)

# Update the weights to their mean before saving
opt.assign_average_vars(model.variables)

model.save('model.h5')

The second option is to update the weight through AverageModelCheckpoint if you set update_weights = True. As the collab notebook example shows

avg_callback = tfa.callbacks.AverageModelCheckpoint(filepath=checkpoint_dir, 
                                                    update_weights=True)
...

#Build Model
model = create_model(moving_avg_sgd)

#Train the network
model.fit(fmnist_train_ds, epochs=5, callbacks=[avg_callback])

Notice that AverageModelCheckpoint also calls assign_average_vars before saving the model, from source code:

def _save_model(self, epoch, logs):
    optimizer = self._get_optimizer()
    assert isinstance(optimizer, AveragedOptimizerWrapper)

    if self.update_weights:
        optimizer.assign_average_vars(self.model.variables)
        return super()._save_model(epoch, logs)
    ...
user_124833460
Related