---
title: StochasticWeightAveraging
description: Implements the Stochastic Weight Averaging (SWA) Callback to average a model.
canonical: >-
  https://lightning.ai/docs/pytorch/stable/optional-api/api/lightning.pytorch.callbacks.StochasticWeightAveraging
---
# StochasticWeightAveraging {#stochasticweightaveraging}

### `class lightning.pytorch.callbacks.StochasticWeightAveraging(swa_lrs, swa_epoch_start=0.8, annealing_epochs=10, annealing_strategy='cos', avg_fn=None, device=device(type='cpu'))` {#lightning.pytorch.callbacks.StochasticWeightAveraging}

Bases: [`Callback`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.callbacks.Callback.md#lightning.pytorch.callbacks.Callback "lightning.pytorch.callbacks.callback.Callback")

Implements the Stochastic Weight Averaging (SWA) Callback to average a model.

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

This documentation is highly inspired by PyTorch’s work on SWA. The callback arguments follow the scheme defined in PyTorch’s `swa_utils` package.

For a SWA explanation, please take a look [here](https://pytorch.org/blog/pytorch-1.6-now-includes-stochastic-weight-averaging).

Warning

This is an [experimental](/docs/pytorch/2.6.6/reference/versioning.md#experimental-api) feature.

Warning

`StochasticWeightAveraging` is currently not supported for multiple optimizers/schedulers.

Warning

`StochasticWeightAveraging` is currently only supported on every epoch.

See also how to [enable it directly on the Trainer](/docs/pytorch/2.6.6/reference/advanced/training_tricks.md#weight-averaging).

Parameters:

*   **swa\_lrs** (`Union`\[[`float`](https://docs.python.org/3/builtins/functions.html#float "(in Python v3.15)"), [`list`](https://docs.python.org/3/builtins/stdtypes.html#list "(in Python v3.15)")\[[`float`](https://docs.python.org/3/builtins/functions.html#float "(in Python v3.15)")\]\]) –
    
    The SWA learning rate to use:
    
    *   `float`. Use this value for all parameter groups of the optimizer.
        
    *   `List[float]`. A list values for each parameter group of the optimizer.
        
    
*   **swa\_epoch\_start** (`Union`\[[`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)"), [`float`](https://docs.python.org/3/builtins/functions.html#float "(in Python v3.15)")\]) – If provided as int, the procedure will start from the `swa_epoch_start`\-th epoch. If provided as float between 0 and 1, the procedure will start from `int(swa_epoch_start * max_epochs)` epoch
    
*   **annealing\_epochs** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – number of epochs in the annealing phase (default: 10)
    
*   **annealing\_strategy** ([`Literal`](https://docs.python.org/3/library/typing.html#typing.Literal "(in Python v3.15)")\[`'cos'`, `'linear'`\]) –
    
    Specifies the annealing strategy (default: “cos”):
    
    *   `"cos"`. For cosine annealing.
        
    *   `"linear"` For linear annealing
        
    
*   **avg\_fn** ([`Optional`](https://docs.python.org/3/library/typing.html#typing.Optional "(in Python v3.15)")\[[`Callable`](https://docs.python.org/3/library/typing.html#typing.Callable "(in Python v3.15)")\[\[[`Tensor`](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "(in PyTorch v2.15)"), [`Tensor`](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "(in PyTorch v2.15)"), [`Tensor`](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "(in PyTorch v2.15)")\], [`Tensor`](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "(in PyTorch v2.15)")\]\]) – the averaging function used to update the parameters; the function must take in the current value of the `AveragedModel` parameter, the current value of `model` parameter and the number of models already averaged; if None, equally weighted average is used (default: `None`)
    
*   **device** (`Union`\[[`device`](https://docs.pytorch.org/docs/stable/tensor_attributes.html#torch.device "(in PyTorch v2.15)"), [`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)"), [`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")\]) – if provided, the averaged model will be stored on the `device`. When None is provided, it will infer the device from `pl_module`. (default: `"cpu"`)
    

#### `static avg_fn(averaged_model_parameter, model_parameter, num_averaged)` {#lightning.pytorch.callbacks.StochasticWeightAveraging.avg_fn}

Adapted from [https://github.com/pytorch/pytorch/blob/v1.7.1/torch/optim/swa\_utils.py#L95-L97](https://github.com/pytorch/pytorch/blob/v1.7.1/torch/optim/swa_utils.py#L95-L97).

Return type:

[`Tensor`](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "(in PyTorch v2.15)")

#### `load_state_dict(state_dict)` {#lightning.pytorch.callbacks.StochasticWeightAveraging.load_state_dict}

Called when loading a checkpoint, implement to reload callback state given callback’s `state_dict`.

Parameters:

**state\_dict** ([`dict`](https://docs.python.org/3/builtins/stdtypes.html#dict "(in Python v3.15)")\[[`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)"), [`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")\]) – the callback state returned by `state_dict`.

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

#### `on_fit_start(trainer, pl_module)` {#lightning.pytorch.callbacks.StochasticWeightAveraging.on_fit_start}

Called when fit begins.

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

#### `on_train_end(trainer, pl_module)` {#lightning.pytorch.callbacks.StochasticWeightAveraging.on_train_end}

Called when the train ends.

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

#### `on_train_epoch_end(trainer, *args)` {#lightning.pytorch.callbacks.StochasticWeightAveraging.on_train_epoch_end}

Called when the train epoch ends.

To access all batch outputs at the end of the epoch, you can cache step outputs as an attribute of the [`lightning.pytorch.core.LightningModule`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.core.LightningModule.md#lightning.pytorch.core.LightningModule "lightning.pytorch.core.LightningModule") and access them in this hook:

```python
class MyLightningModule(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.training_step_outputs = []

    def training_step(self):
        loss = ...
        self.training_step_outputs.append(loss)
        return loss


class MyCallback(L.Callback):
    def on_train_epoch_end(self, trainer, pl_module):
        # do something with all training_step outputs, for example:
        epoch_mean = torch.stack(pl_module.training_step_outputs).mean()
        pl_module.log("training_epoch_mean", epoch_mean)
        # free up the memory
        pl_module.training_step_outputs.clear()
```

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

#### `on_train_epoch_start(trainer, pl_module)` {#lightning.pytorch.callbacks.StochasticWeightAveraging.on_train_epoch_start}

Called when the train epoch begins.

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

#### `reset_batch_norm_and_save_state(pl_module)` {#lightning.pytorch.callbacks.StochasticWeightAveraging.reset_batch_norm_and_save_state}

Adapted from [https://github.com/pytorch/pytorch/blob/v1.7.1/torch/optim/swa\_utils.py#L140-L154](https://github.com/pytorch/pytorch/blob/v1.7.1/torch/optim/swa_utils.py#L140-L154).

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

#### `reset_momenta()` {#lightning.pytorch.callbacks.StochasticWeightAveraging.reset_momenta}

Adapted from [https://github.com/pytorch/pytorch/blob/v1.7.1/torch/optim/swa\_utils.py#L164-L165](https://github.com/pytorch/pytorch/blob/v1.7.1/torch/optim/swa_utils.py#L164-L165).

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

#### `setup(trainer, pl_module, stage)` {#lightning.pytorch.callbacks.StochasticWeightAveraging.setup}

Called when fit, validate, test, predict, or tune begins.

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

#### `state_dict()` {#lightning.pytorch.callbacks.StochasticWeightAveraging.state_dict}

Called when saving a checkpoint, implement to generate callback’s `state_dict`.

Return type:

[`dict`](https://docs.python.org/3/builtins/stdtypes.html#dict "(in Python v3.15)")\[[`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)"), [`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")\]

Returns:

A dictionary containing callback state.

#### `static update_parameters(average_model, model, n_averaged, avg_fn)` {#lightning.pytorch.callbacks.StochasticWeightAveraging.update_parameters}

Adapted from [https://github.com/pytorch/pytorch/blob/v1.7.1/torch/optim/swa\_utils.py#L104-L112](https://github.com/pytorch/pytorch/blob/v1.7.1/torch/optim/swa_utils.py#L104-L112).

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

---

**Navigation**

- Parent: [Api](/docs/pytorch/2.6.6/optional-api/api.md)