Optimization¶
The module pyro.optim provides support for optimization in Pyro. In particular it provides PyroOptim, which is used to wrap PyTorch optimizers and manage optimizers for dynamically generated parameters (see the tutorial SVI Part I for a discussion). Any custom optimization algorithms are also to be found here.
Pyro Optimizers¶
-
class
PyroOptim
(optim_constructor, optim_args)[source]¶ Bases:
object
A wrapper for torch.optim.Optimizer objects that helps with managing dynamically generated parameters.
Parameters: - optim_constructor – a torch.optim.Optimizer
- optim_args – a dictionary of learning arguments for the optimizer or a callable that returns such dictionaries
-
__call__
(params, *args, **kwargs)[source]¶ Parameters: params (an iterable of strings) – a list of parameters Do an optimization step for each param in params. If a given param has never been seen before, initialize an optimizer for it.
-
get_state
()[source]¶ Get state associated with all the optimizers in the form of a dictionary with key-value pairs (parameter name, optim state dicts)
-
AdagradRMSProp
(optim_args)[source]¶ A wrapper for an optimizer that is a mash-up of
Adagrad
andRMSprop
.
-
ClippedAdam
(optim_args)[source]¶ A wrapper for a modification of the
Adam
optimization algorithm that supports gradient clipping.
-
class
PyroLRScheduler
(scheduler_constructor, optim_args)[source]¶ Bases:
pyro.optim.optim.PyroOptim
A wrapper for torch.optim.lr_scheduler objects that adjust learning rates for dynamically generated parameters.
Parameters: - optim_constructor – a torch.optim.lr_scheduler
- optim_args – a dictionary of learning arguments for the optimizer or a callable that returns such dictionaries. must contain the key ‘optimizer’ with pytorch optimizer value
Example:
optimizer = torch.optim.SGD pyro_scheduler = pyro.optim.ExponentialLR({'optimizer': optimizer, 'optim_args': {'lr': 0.01}, 'gamma': 0.1}) svi = SVI(model, guide, pyro_scheduler, loss=TraceGraph_ELBO()) svi.step()
PyTorch Optimizers¶
-
Adadelta
(optim_args)¶ Wraps
torch.optim.Adadelta
withPyroOptim
.
-
Adagrad
(optim_args)¶ Wraps
torch.optim.Adagrad
withPyroOptim
.
-
Adam
(optim_args)¶ Wraps
torch.optim.Adam
withPyroOptim
.
-
AdamW
(optim_args)¶ Wraps
torch.optim.AdamW
withPyroOptim
.
-
SparseAdam
(optim_args)¶ Wraps
torch.optim.SparseAdam
withPyroOptim
.
-
Adamax
(optim_args)¶ Wraps
torch.optim.Adamax
withPyroOptim
.
-
ASGD
(optim_args)¶ Wraps
torch.optim.ASGD
withPyroOptim
.
-
SGD
(optim_args)¶ Wraps
torch.optim.SGD
withPyroOptim
.
-
Rprop
(optim_args)¶ Wraps
torch.optim.Rprop
withPyroOptim
.
-
RMSprop
(optim_args)¶ Wraps
torch.optim.RMSprop
withPyroOptim
.
-
LambdaLR
(optim_args)¶ Wraps
torch.optim.LambdaLR
withPyroLRScheduler
.
-
StepLR
(optim_args)¶ Wraps
torch.optim.StepLR
withPyroLRScheduler
.
-
MultiStepLR
(optim_args)¶ Wraps
torch.optim.MultiStepLR
withPyroLRScheduler
.
-
ExponentialLR
(optim_args)¶ Wraps
torch.optim.ExponentialLR
withPyroLRScheduler
.
-
CosineAnnealingLR
(optim_args)¶ Wraps
torch.optim.CosineAnnealingLR
withPyroLRScheduler
.
-
ReduceLROnPlateau
(optim_args)¶ Wraps
torch.optim.ReduceLROnPlateau
withPyroLRScheduler
.
-
CyclicLR
(optim_args)¶ Wraps
torch.optim.CyclicLR
withPyroLRScheduler
.
-
CosineAnnealingWarmRestarts
(optim_args)¶ Wraps
torch.optim.CosineAnnealingWarmRestarts
withPyroLRScheduler
.
Higher-Order Optimizers¶
-
class
MultiOptimizer
[source]¶ Bases:
object
Base class of optimizers that make use of higher-order derivatives.
Higher-order optimizers generally use
torch.autograd.grad()
rather thantorch.Tensor.backward()
, and therefore require a different interface from usual Pyro and PyTorch optimizers. In this interface, thestep()
method inputs aloss
tensor to be differentiated, and backpropagation is triggered one or more times inside the optimizer.Derived classes must implement
step()
to compute derivatives and update parameters in-place.Example:
tr = poutine.trace(model).get_trace(*args, **kwargs) loss = -tr.log_prob_sum() params = {name: site['value'].unconstrained() for name, site in tr.nodes.items() if site['type'] == 'param'} optim.step(loss, params)
-
step
(loss, params)[source]¶ Performs an in-place optimization step on parameters given a differentiable
loss
tensor.Note that this detaches the updated tensors.
Parameters: - loss (torch.Tensor) – A differentiable tensor to be minimized. Some optimizers require this to be differentiable multiple times.
- params (dict) – A dictionary mapping param name to unconstrained value as stored in the param store.
-
get_step
(loss, params)[source]¶ Computes an optimization step of parameters given a differentiable
loss
tensor, returning the updated values.Note that this preserves derivatives on the updated tensors.
Parameters: - loss (torch.Tensor) – A differentiable tensor to be minimized. Some optimizers require this to be differentiable multiple times.
- params (dict) – A dictionary mapping param name to unconstrained value as stored in the param store.
Returns: A dictionary mapping param name to updated unconstrained value.
Return type:
-
-
class
PyroMultiOptimizer
(optim)[source]¶ Bases:
pyro.optim.multi.MultiOptimizer
Facade to wrap
PyroOptim
objects in aMultiOptimizer
interface.
-
class
TorchMultiOptimizer
(optim_constructor, optim_args)[source]¶ Bases:
pyro.optim.multi.PyroMultiOptimizer
Facade to wrap
Optimizer
objects in aMultiOptimizer
interface.
-
class
MixedMultiOptimizer
(parts)[source]¶ Bases:
pyro.optim.multi.MultiOptimizer
Container class to combine different
MultiOptimizer
instances for different parameters.Parameters: parts (list) – A list of (names, optim)
pairs, where eachnames
is a list of parameter names, and eachoptim
is aMultiOptimizer
orPyroOptim
object to be used for the named parameters. Together thenames
should partition up all desired parameters to optimize.Raises: ValueError – if any name is optimized by multiple optimizers.
-
class
Newton
(trust_radii={})[source]¶ Bases:
pyro.optim.multi.MultiOptimizer
Implementation of
MultiOptimizer
that performs a Newton update on batched low-dimensional variables, optionally regularizing via a per-parametertrust_radius
. Seenewton_step()
for details.The result of
get_step()
will be differentiable, however the updated values fromstep()
will be detached.Parameters: trust_radii (dict) – a dict mapping parameter name to radius of trust region. Missing names will use unregularized Newton update, equivalent to infinite trust radius.