StochasticWeightAverage

class mmengine.model.StochasticWeightAverage(model, interval=1, device=None, update_buffers=False)

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 (torch.nn.modules.module.Module) –

• interval (int) –

• device (Optional[torch.device]) –

• update_buffers (bool) –

Return type

None

avg_func(averaged_param, source_param, steps)

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