del on a given dataloader ` loader  ` at the end of
training:

    
    
    >>> torch.optim.swa_utils.update_bn(loader, swa_model)
    

` update_bn()  ` applies the ` swa_model  ` to every element in the dataloader
and computes the activation statistics for each batch normalization layer in
the model.

Warning

` update_bn()  ` assumes that each batch in the dataloader ` loader  ` is
either a tensors or a list of tensors where the first element is the tensor
that the network ` swa_model  ` should be applied to. If your dataloader has a
different structure, you can update the batch normalization statistics of the
` swa_model  ` by doing a forward pass with the ` swa_model  ` on each element
of the dataset.

###  Custom averaging strategies  ¶

By default, ` torch.optim.swa_utils.AveragedModel  ` computes a running equal
average of the parameters that you provide, but you can also use custom
averaging functions with the ` avg_fn  ` parameter. In the following example `
ema_model  ` computes an exponential moving average.

Example:

    
    
    >>> ema_avg = lambda averaged_model_parameter, model_parameter, num_averaged:\
    >>>         0.1 * averaged_model_parameter + 0.9 * model_parameter
    >>> ema_model = torch.optim.swa_utils.AveragedModel(model, avg_fn=ema_avg)
    

###  Putting it all together  ¶

In the example below, ` swa_model  ` is the SWA model that accumulates the
averages of the weights. We train the model for a total of 300 epochs and we
switch to the SWA learning rate schedule and start to collect SWA averages of
the parameters at epoch 160:

    
    
    >>> loader, optimizer, model, loss_fn = ...
    >>> swa_model = torch.optim.swa_utils.AveragedModel(model)
    >>> scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=300)
    >>> swa_start = 160
    >>> swa_scheduler = SWALR(optimizer, swa_lr=0.05)
    >>>
    >>> for epoch in range(300):
    >>>       for input, target in loader:
    >>>           optimizer.zero_grad()
    >>>           loss_fn(model(input), target).backward()
    >>>           optimizer.step()
    >>>       if i > swa_start:
    >>>           swa_model.update_parameters(model)
    >>>           swa_scheduler.step()
    >>>       else:
    >>>           scheduler.step()
    >>>
    >>> # Update bn statistics for the swa_model at the end
    >>> torch.optim.swa_utils.update_bn(loader, swa_model)
    >>> # Use swa_model to make predictions on test data
    >>> preds = swa_model(test_input)
    

[ Next ![](_static/images/chevron-right-orange.svg) ](complex_numbers.html
"Complex Numbers") [ ![](_static/images/chevron-right-orange.svg) Previous
](onnx.html "torch.onnx")

* * *

© Copyright 2019, Torch Contributors.

Built with [ Sphinx ](http://sphinx-doc.org/) using a [ theme
](https://github.com/rtfd/sphinx_rtd_theme) provided by [ Read the Docs
](https://readthedocs.org) .

  * torch.optim 
    * How to use an optimizer 
      * Constructing it 
      * Per-parameter options 
      * Taking an optimization step 
        * ` optimizer.step()  `
        * ` optimizer.step(closure)  `
    * Algorithms 
    * How to adjust learning rate 
    * Stochastic Weight Averaging 
      * Constructing averaged models 
      * SWA learning rate schedules 
      * Taking care of batch normalization 
      * Custom averaging strategies 
      * Putting it all together 

![](https://www.googleadservices.com/pagead/conversion/795629140/?label=txkmCPmdtosBENSssfsC&guid=ON&script=0)

##  Docs

Access comprehensive developer documentation for PyTorch

[ View Docs ](https://pytorch.org/docs/stable/index.html)

##  Tutorials

Get in-depth tutorials for beginners and advanced developers

[ View Tutorials ](https://pytorch.org/tutorials)

##  Resources

Find development resources and get your questions answered

[ View Resources ](https://pytorch.org/resources)

[ ](https://pytorch.org/)

  * [ PyTorch ](https://pytorch.org/)
  * [ Get Started ](https://pytorch.org/get-started)
  * [ Feature