StochasticWeightAverage — mmengine 0.10.7 documentation (original) (raw)

class mmengine.model.StochasticWeightAverage(model, interval=1, device=None, update_buffers=False)[source]¶

Implements the stochastic weight averaging (SWA) of the model.

Stochastic Weight Averaging was proposed in Averaging Weights Leads to Wider Optima and Better Generalization, UAI 2018. by Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov and Andrew Gordon Wilson.

Parameters:

• model (Module) –
• interval (int) –
• device (device | None) –
• update_buffers (bool) –

avg_func(averaged_param, source_param, steps)[source]¶

Compute the average of the parameters using stochastic weight average.

Parameters:

• averaged_param (Tensor) – The averaged parameters.
• source_param (Tensor) – The source parameters.
• steps (int) – The number of times the parameters have been updated.

Return type:

None