View source on GitHub 
This class extends optimizers with Stochastic Weight Averaging (SWA).
Inherits From: AveragedOptimizerWrapper
tfa.optimizers.SWA(
optimizer: tfa.types.Optimizer
,
start_averaging: int = 0,
average_period: int = 10,
name: str = 'SWA',
**kwargs
)
Used in the notebooks
Used in the tutorials 

The Stochastic Weight Averaging mechanism was proposed by Pavel Izmailov
et. al in the paper Averaging Weights Leads to Wider Optima and
Better Generalization. The optimizer
implements averaging of multiple points along the trajectory of SGD. The
optimizer expects an inner optimizer which will be used to apply the
gradients to the variables and itself computes a running average of the
variables every k
steps (which generally corresponds to the end
of a cycle when a cyclic learning rate is employed).
We also allow the specification of the number of steps averaging
should first happen after. Let's say, we want averaging to happen every k
steps after the first m
steps. After step m
we'd take a snapshot of the
variables and then average the weights appropriately at step m + k
,
m + 2k
and so on. The assign_average_vars function can be called at the
end of training to obtain the averaged_weights from the optimizer.
Example of usage:
opt = tf.keras.optimizers.SGD(learning_rate)
opt = tfa.optimizers.SWA(opt, start_averaging=m, average_period=k)
Args  

optimizer

The original optimizer that will be used to compute and apply the gradients. 
start_averaging

An integer. Threshold to start averaging using
SWA. Averaging only occurs at start_averaging iters, must
be >= 0. If start_averaging = m, the first snapshot will be
taken after the mth application of gradients (where the first
iteration is iteration 0).

average_period

An integer. The synchronization period of SWA. The averaging occurs every average_period steps. Averaging period needs to be >= 1. 
name

Optional name for the operations created when applying gradients. Defaults to 'SWA'. 
**kwargs

keyword arguments. Allowed to be {clipnorm ,
clipvalue , lr , decay }. clipnorm is clip gradients by
norm; clipvalue is clip gradients by value, decay is
included for backward compatibility to allow time inverse
decay of learning rate. lr is included for backward
compatibility, recommended to use learning_rate instead.

Attributes  

clipnorm

float or None . If set, clips gradients to a maximum norm.

clipvalue

float or None . If set, clips gradients to a maximum value.

global_clipnorm

float or None .
If set, clips gradients to a maximum norm. Check 
iterations

Variable. The number of training steps this Optimizer has run. 
learning_rate


lr


weights

Returns variables of this Optimizer based on the order created. 
Methods
add_slot
add_slot(
var, slot_name, initializer='zeros', shape=None
)
Add a new slot variable for var
.
A slot variable is an additional variable associated with var
to train.
It is allocated and managed by optimizers, e.g. Adam
.
Args  

var

a Variable object.

slot_name

name of the slot variable. 
initializer

initializer of the slot variable 
shape

(Optional) shape of the slot variable. If not set, it will default
to the shape of var .

Returns  

A slot variable. 
add_weight
add_weight(
name,
shape,
dtype=None,
initializer='zeros',
trainable=None,
synchronization=tf.VariableSynchronization.AUTO,
aggregation=tf.VariableAggregation.NONE
)
apply_gradients
apply_gradients(
grads_and_vars, name=None, **kwargs
)
Apply gradients to variables.
This is the second part of minimize()
. It returns an Operation
that
applies gradients.
The method sums gradients from all replicas in the presence of
tf.distribute.Strategy
by default. You can aggregate gradients yourself by
passing experimental_aggregate_gradients=False
.
Example:
grads = tape.gradient(loss, vars)
grads = tf.distribute.get_replica_context().all_reduce('sum', grads)
# Processing aggregated gradients.
optimizer.apply_gradients(zip(grads, vars),
experimental_aggregate_gradients=False)
Args  

grads_and_vars

List of (gradient, variable) pairs. 
name

Optional name for the returned operation. Default to the name passed
to the Optimizer constructor.

experimental_aggregate_gradients

Whether to sum gradients from different
replicas in the presence of tf.distribute.Strategy . If False, it's
user responsibility to aggregate the gradients. Default to True.

Returns  

An Operation that applies the specified gradients. The iterations
will be automatically increased by 1.

Raises  

TypeError

If grads_and_vars is malformed.

ValueError

If none of the variables have gradients. 
RuntimeError

If called in a crossreplica context. 
assign_average_vars
assign_average_vars(
var_list
)
Assign variables in var_list with their respective averages.
Args  

var_list

List of model variables to be assigned to their average. 
Returns  

assign_op

The op corresponding to the assignment operation of variables to their average. 
Example:
model = tf.Sequential([...])
opt = tfa.optimizers.SWA(
tf.keras.optimizers.SGD(lr=2.0), 100, 10)
model.compile(opt, ...)
model.fit(x, y, ...)
# Update the weights to their mean before saving
opt.assign_average_vars(model.variables)
model.save('model.h5')
average_op
@tf.function
average_op( var, average_var, local_apply_state )
from_config
@classmethod
from_config( config, custom_objects=None )
Creates an optimizer from its config.
This method is the reverse of get_config
,
capable of instantiating the same optimizer from the config
dictionary.
Args  

config

A Python dictionary, typically the output of get_config. 
custom_objects

A Python dictionary mapping names to additional Python objects used to create this optimizer, such as a function used for a hyperparameter. 
Returns  

An optimizer instance. 
get_config
get_config()
Returns the config of the optimizer.
An optimizer config is a Python dictionary (serializable) containing the configuration of an optimizer. The same optimizer can be reinstantiated later (without any saved state) from this configuration.
Returns  

Python dictionary. 
get_gradients
get_gradients(
loss, params
)
Returns gradients of loss
with respect to params
.
Should be used only in legacy v1 graph mode.
Args  

loss

Loss tensor. 
params

List of variables. 
Returns  

List of gradient tensors. 
Raises  

ValueError

In case any gradient cannot be computed (e.g. if gradient function not implemented). 
get_slot
get_slot(
var, slot_name
)
get_slot_names
get_slot_names()
A list of names for this optimizer's slots.
get_updates
get_updates(
loss, params
)
get_weights
get_weights()
Returns the current weights of the optimizer.
The weights of an optimizer are its state (ie, variables). This function returns the weight values associated with this optimizer as a list of Numpy arrays. The first value is always the iterations count of the optimizer, followed by the optimizer's state variables in the order they were created. The returned list can in turn be used to load state into similarly parameterized optimizers.
For example, the RMSprop optimizer for this simple model returns a list of three values the iteration count, followed by the rootmeansquare value of the kernel and bias of the single Dense layer:
opt = tf.keras.optimizers.RMSprop()
m = tf.keras.models.Sequential([tf.keras.layers.Dense(10)])
m.compile(opt, loss='mse')
data = np.arange(100).reshape(5, 20)
labels = np.zeros(5)
results = m.fit(data, labels) # Training.
len(opt.get_weights())
3
Returns  

Weights values as a list of numpy arrays. 
minimize
minimize(
loss, var_list, grad_loss=None, name=None, tape=None
)
Minimize loss
by updating var_list
.
This method simply computes gradient using tf.GradientTape
and calls
apply_gradients()
. If you want to process the gradient before applying
then call tf.GradientTape
and apply_gradients()
explicitly instead
of using this function.
Args  

loss

Tensor or callable. If a callable, loss should take no arguments
and return the value to minimize. If a Tensor , the tape argument
must be passed.

var_list

list or tuple of Variable objects to update to minimize
loss , or a callable returning the list or tuple of Variable objects.
Use callable when the variable list would otherwise be incomplete before
minimize since the variables are created at the first time loss is
called.

grad_loss

(Optional). A Tensor holding the gradient computed for
loss .

name

(Optional) str. Name for the returned operation. 
tape

(Optional) tf.GradientTape . If loss is provided as a Tensor ,
the tape that computed the loss must be provided.

Returns  

An Operation that updates the variables in var_list . The iterations
will be automatically increased by 1.

Raises  

ValueError

If some of the variables are not Variable objects.

set_weights
set_weights(
weights
)
Set the weights of the optimizer.
The weights of an optimizer are its state (ie, variables). This function takes the weight values associated with this optimizer as a list of Numpy arrays. The first value is always the iterations count of the optimizer, followed by the optimizer's state variables in the order they are created. The passed values are used to set the new state of the optimizer.
For example, the RMSprop optimizer for this simple model takes a list of three values the iteration count, followed by the rootmeansquare value of the kernel and bias of the single Dense layer:
opt = tf.keras.optimizers.RMSprop()
m = tf.keras.models.Sequential([tf.keras.layers.Dense(10)])
m.compile(opt, loss='mse')
data = np.arange(100).reshape(5, 20)
labels = np.zeros(5)
results = m.fit(data, labels) # Training.
new_weights = [np.array(10), np.ones([20, 10]), np.zeros([10])]
opt.set_weights(new_weights)
opt.iterations
<tf.Variable 'RMSprop/iter:0' shape=() dtype=int64, numpy=10>
Args  

weights

weight values as a list of numpy arrays. 
variables
variables()
Returns variables of this Optimizer based on the order created.