Constructing it#

To construct anOptimizer you have to give it an iterable containing theparameters (all should beParameter s) or named parameters(tuples of (str,Parameter)) to optimize. Then,you can specify optimizer-specific options such as the learning rate, weight decay, etc.

Example:

optimizer=optim.SGD(model.parameters(),lr=0.01,momentum=0.9)optimizer=optim.Adam([var1,var2],lr=0.0001)

Named parameters example:

optimizer=optim.SGD(model.named_parameters(),lr=0.01,momentum=0.9)optimizer=optim.Adam([('layer0',var1),('layer1',var2)],lr=0.0001)

Per-parameter options#

Optimizer s also support specifying per-parameter options. To do this, insteadof passing an iterable ofVariable s, pass in an iterable ofdict s. Each of them will define a separate parameter group, and should containaparams key, containing a list of parameters belonging to it. Other keysshould match the keyword arguments accepted by the optimizers, and will be usedas optimization options for this group.

For example, this is very useful when one wants to specify per-layer learning rates:

optim.SGD([{'params':model.base.parameters(),'lr':1e-2},{'params':model.classifier.parameters()}],lr=1e-3,momentum=0.9)optim.SGD([{'params':model.base.named_parameters(),'lr':1e-2},{'params':model.classifier.named_parameters()}],lr=1e-3,momentum=0.9)

This means thatmodel.base’s parameters will use a learning rate of1e-2, whereasmodel.classifier’s parameters will stick to the default learning rate of1e-3.Finally a momentum of0.9 will be used for all parameters.

Also consider the following example related to the distinct penalization of parameters.Remember thatparameters() returns an iterable thatcontains all learnable parameters, including biases and otherparameters that may prefer distinct penalization. To address this, one can specifyindividual penalization weights for each parameter group:

bias_params=[pforname,pinself.named_parameters()if'bias'inname]others=[pforname,pinself.named_parameters()if'bias'notinname]optim.SGD([{'params':others},{'params':bias_params,'weight_decay':0}],weight_decay=1e-2,lr=1e-2)

In this manner, bias terms are isolated from non-bias terms, and aweight_decayof0 is set specifically for the bias terms, as to avoid any penalization forthis group.

Taking an optimization step#

All optimizers implement astep() method, that updates theparameters. It can be used in two ways:

forinput,targetindataset:optimizer.zero_grad()output=model(input)loss=loss_fn(output,target)loss.backward()optimizer.step()

forinput,targetindataset:defclosure():optimizer.zero_grad()output=model(input)loss=loss_fn(output,target)loss.backward()returnlossoptimizer.step(closure)

Warning

Prior to PyTorch 1.1.0, the learning rate scheduler was expected to be called beforethe optimizer’s update; 1.1.0 changed this behavior in a BC-breaking way. If you usethe learning rate scheduler (callingscheduler.step()) before the optimizer’s update(callingoptimizer.step()), this will skip the first value of the learning rate schedule.If you are unable to reproduce results after upgrading to PyTorch 1.1.0, please checkif you are callingscheduler.step() at the wrong time.

Constructing averaged models#

TheAveragedModel class serves to compute the weights of the SWA or EMA model.

You can create an SWA averaged model by running:

>>>averaged_model=AveragedModel(model)

EMA models are constructed by specifying themulti_avg_fn argument as follows:

>>>decay=0.999>>>averaged_model=AveragedModel(model,multi_avg_fn=get_ema_multi_avg_fn(decay))

Decay is a parameter between 0 and 1 that controls how fast the averaged parameters are decayed. If not provided totorch.optim.swa_utils.get_ema_multi_avg_fn(), the default is 0.999. Decay value should be close to 1.0, as smaller values can cause optimization convergence issues.

torch.optim.swa_utils.get_ema_multi_avg_fn() returns a function that applies the following EMA equation to the weights:

where alpha is the EMA decay.

Here the modelmodel can be an arbitrarytorch.nn.Module object.averaged_modelwill keep track of the running averages of the parameters of themodel. To update theseaverages, you should use theupdate_parameters() function after theoptimizer.step():

>>>averaged_model.update_parameters(model)

For SWA and EMA, this call is usually done right after the optimizerstep(). In the case of SWA, this is usually skipped for some numbers of steps at the beginning of the training.

Custom averaging strategies#

By default,torch.optim.swa_utils.AveragedModel computes a running equal average ofthe parameters that you provide, but you can also use custom averaging functions with theavg_fn ormulti_avg_fn parameters:

• avg_fn allows defining a function operating on each parameter tuple (averaged parameter, model parameter) and should return the new averaged parameter.

• multi_avg_fn allows defining more efficient operations acting on a tuple of parameter lists, (averaged parameter list, model parameter list), at the same time, for example using thetorch._foreach* functions. This function must update the averaged parameters in-place.

In the following exampleema_model computes an exponential moving average using theavg_fn parameter:

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

In the following exampleema_model computes an exponential moving average using the more efficientmulti_avg_fn parameter:

>>>ema_model=AveragedModel(model,multi_avg_fn=get_ema_multi_avg_fn(0.9))

SWA learning rate schedules#

Typically, in SWA the learning rate is set to a high constant value.SWALR is alearning rate scheduler that anneals the learning rate to a fixed value, and then keeps itconstant. For example, the following code creates a scheduler that linearly anneals thelearning rate from its initial value to 0.05 in 5 epochs within each parameter group:

>>>swa_scheduler=torch.optim.swa_utils.SWALR(optimizer, \>>>anneal_strategy="linear",anneal_epochs=5,swa_lr=0.05)

You can also use cosine annealing to a fixed value instead of linear annealing by settinganneal_strategy="cos".

Taking care of batch normalization#

update_bn() is a utility function that allows to compute the batchnorm statistics for the SWA modelon a given dataloaderloader at the end of training:

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

update_bn() applies theswa_model to every element in the dataloader and computes the activationstatistics for each batch normalization layer in the model.

Putting it all together: SWA#

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 scheduleand 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)>>>>>>forepochinrange(300):>>>forinput,targetinloader:>>>optimizer.zero_grad()>>>loss_fn(model(input),target).backward()>>>optimizer.step()>>>ifepoch>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)

Putting it all together: EMA#

In the example below,ema_model is the EMA model that accumulates the exponentially-decayed averages of the weights with a decay rate of 0.999.We train the model for a total of 300 epochs and start to collect EMA averages immediately.

>>>loader,optimizer,model,loss_fn=...>>>ema_model=torch.optim.swa_utils.AveragedModel(model, \>>>multi_avg_fn=torch.optim.swa_utils.get_ema_multi_avg_fn(0.999))>>>>>>forepochinrange(300):>>>forinput,targetinloader:>>>optimizer.zero_grad()>>>loss_fn(model(input),target).backward()>>>optimizer.step()>>>ema_model.update_parameters(model)>>>>>># Update bn statistics for the ema_model at the end>>>torch.optim.swa_utils.update_bn(loader,ema_model)>>># Use ema_model to make predictions on test data>>>preds=ema_model(test_input)

torch.optim.swa_utils.get_ema_multi_avg_fn(decay=0.999)[source]#

Get the function applying exponential moving average (EMA) across multiple params.

torch.optim.swa_utils.update_bn(loader,model,device=None)[source]#

Update BatchNorm running_mean, running_var buffers in the model.

It performs one pass over data inloader to estimate the activationstatistics for BatchNorm layers in the model.

Parameters

• loader (torch.utils.data.DataLoader) – dataset loader to compute theactivation statistics on. Each data batch should be either atensor, or a list/tuple whose first element is a tensorcontaining data.

• model (torch.nn.Module) – model for which we seek to update BatchNormstatistics.

• device (torch.device,optional) – If set, data will be transferred todevice before being passed intomodel.

Example

>>>loader,model=...>>>torch.optim.swa_utils.update_bn(loader,model)

Note

Theupdate_bn utility assumes that each data batch inloaderis either a tensor or a list or tuple of tensors; in the latter case itis assumed thatmodel.forward() should be called on the firstelement of the list or tuple corresponding to the data batch.