Module: tfa.optimizers.stochastic_weight_averaging

View source on GitHub

An implementation of the Stochastic Weight Averaging optimizer.

The Stochastic Weight Averaging mechanism was proposed by Pavel Izmailov et. al in the paper Averaging Weights Leads to Wider Optima and Better Generalization. The optimizer implements averaging of multiple points along the trajectory of SGD. This averaging has shown to improve model performance on validation/test sets whilst possibly causing a small increase in loss on the training set.


class SWA: This class extends optimizers with Stochastic Weight Averaging (SWA).