kristpapadopoulos / keras-stochastic-weight-averaging

Keras callback function for stochastic weight averaging

Stochastic Weight Averaging with Keras callback function

Stochastic Weight Averaging following paper Averaging Weights Leads to Wider Optima and Better Generalization

The file swa.py contains an implementation for stochastic weight averaging (SWA) with a constant learning rate for a user defined amount of epochs.

Callback is instantiated with filename for saving the final weights of the model after SWA and the number of epochs to average.

Example

The total number of training epochs 150, SWA to start from epoch 140 to average last 10 epochs.

from swa import SWA

# specify number of training epochs
number_of_epochs = 150

# specify the start epoch of stochastic weight averaging
swa = SWA(140, filepath = None)

# call SWA during model fitting
model.fit(..., callbacks = [swa])

keras-stochastic-weight-averaging's Issues

should we do something special for BN layers?

because in the swa (pytorch) https://github.com/timgaripov/swa
there is one step that use the data to update the BN layer? like
utils.bn_update(loaders['train'], swa_model)
is it necessary?

Which optimizer should I use?

Thanks for your code, but I am curious that which optimizer should I use? Can I use adam? Or I must use SGD? Thank you.

Would you like to give us an example to use swa callback?

Thank you so much for your code.
It will be great if you would like to upload an example to show how to use your callback. If I use learning_rate_schedular, may I combine your saw together? Because you said that your learning rate is a constant. Many thanks.