5561

pending completion

Build # 5561

Build Type

Pull #2191

travis-ci-com

Committed by

web-flow

Commit Message

asymptotic complexity comment

- addresses https://github.com/pyro-ppl/pyro/pull/2191/files/14717a3bb#r352402571

Pull Request Pull Request #2191: Reweighted Wake-Sleep

Run Details

113 of 113 new or added lines in 5 files covered. (100.0%)

29756 of 40587 relevant lines covered (73.31%)

2.21 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

66.07

/pyro/infer/svi.py

import warnings

import torch

import pyro
import pyro.optim
import pyro.poutine as poutine
from pyro.infer.abstract_infer import TracePosterior
from pyro.infer.elbo import ELBO
from pyro.infer.util import torch_item


class SVI(TracePosterior):
    """
    :param model: the model (callable containing Pyro primitives)
    :param guide: the guide (callable containing Pyro primitives)
    :param optim: a wrapper a for a PyTorch optimizer
    :type optim: pyro.optim.PyroOptim
    :param loss: an instance of a subclass of :class:`~pyro.infer.elbo.ELBO`.
        Pyro provides three built-in losses:
        :class:`~pyro.infer.trace_elbo.Trace_ELBO`,
        :class:`~pyro.infer.tracegraph_elbo.TraceGraph_ELBO`, and
        :class:`~pyro.infer.traceenum_elbo.TraceEnum_ELBO`.
        See the :class:`~pyro.infer.elbo.ELBO` docs to learn how to implement
        a custom loss.
    :type loss: pyro.infer.elbo.ELBO
    :param num_samples: (DEPRECATED) the number of samples for Monte Carlo posterior approximation
    :param num_steps: (DEPRECATED) the number of optimization steps to take in ``run()``

    A unified interface for stochastic variational inference in Pyro. The most
    commonly used loss is ``loss=Trace_ELBO()``. See the tutorial
    `SVI Part I <http://pyro.ai/examples/svi_part_i.html>`_ for a discussion.
    """
    def __init__(self,
                 model,
                 guide,
                 optim,
                 loss,
                 loss_and_grads=None,
                 num_samples=0,
                 num_steps=0,
                 **kwargs):
        if num_steps:
            warnings.warn('The `num_steps` argument to SVI is deprecated and will be removed in '
                          'a future release. Use `SVI.step` directly to control the '
                          'number of iterations.', FutureWarning)
        if num_samples:
            warnings.warn('The `num_samples` argument to SVI is deprecated and will be removed in '
                          'a future release. Use `pyro.infer.Predictive` class to draw '
                          'samples from the posterior.', FutureWarning)

        self.model = model
        self.guide = guide
        self.optim = optim
        self.num_steps = num_steps
        self.num_samples = num_samples
        super(SVI, self).__init__(**kwargs)

        if not isinstance(optim, pyro.optim.PyroOptim):
            raise ValueError("Optimizer should be an instance of pyro.optim.PyroOptim class.")

        if isinstance(loss, ELBO):
            self.loss = loss.loss
            self.loss_and_grads = loss.loss_and_grads
        else:
            if loss_and_grads is None:
                def _loss_and_grads(*args, **kwargs):
                    loss_val = loss(*args, **kwargs)
                    loss_val.backward(retain_graph=True)
                    return loss_val
                loss_and_grads = _loss_and_grads
            self.loss = loss
            self.loss_and_grads = loss_and_grads

    def run(self, *args, **kwargs):
        """
        .. warning::
            This method is deprecated, and will be removed in a future release.
            For inference, use :meth:`step` directly, and for predictions,
            use the :class:`~pyro.infer.predictive.Predictive` class.
        """
        warnings.warn('The `SVI.run` method is deprecated and will be removed in a '
                      'future release. For inference, use `SVI.step` directly, '
                      'and for predictions, use the `pyro.infer.Predictive` class.',
                      FutureWarning)
        if self.num_steps > 0:
            with poutine.block():
                for i in range(self.num_steps):
                    self.step(*args, **kwargs)
        return super(SVI, self).run(*args, **kwargs)

    def _traces(self, *args, **kwargs):
        for i in range(self.num_samples):
            guide_trace = poutine.trace(self.guide).get_trace(*args, **kwargs)
            model_trace = poutine.trace(poutine.replay(self.model, trace=guide_trace)).get_trace(*args, **kwargs)
            yield model_trace, 1.0

    def evaluate_loss(self, *args, **kwargs):
        """
        :returns: estimate of the loss
        :rtype: float

        Evaluate the loss function. Any args or kwargs are passed to the model and guide.
        """
        with torch.no_grad():
            loss = self.loss(self.model, self.guide, *args, **kwargs)
            if isinstance(loss, tuple):
                return torch_item(loss[0]), torch_item(loss[1])
            else:
                return torch_item(loss)

    def step(self, *args, **kwargs):
        """
        :returns: estimate of the loss
        :rtype: float

        Take a gradient step on the loss function (and any auxiliary loss functions
        generated under the hood by `loss_and_grads`).
        Any args or kwargs are passed to the model and guide
        """
        # get loss and compute gradients
        with poutine.trace(param_only=True) as param_capture:
            loss = self.loss_and_grads(self.model, self.guide, *args, **kwargs)

        params = set(site["value"].unconstrained()
                     for site in param_capture.trace.nodes.values())

        # actually perform gradient steps
        # torch.optim objects gets instantiated for any params that haven't been seen yet
        self.optim(params)

        # zero gradients
        pyro.infer.util.zero_grads(params)

        # RWS: handle loss = (wake_theta_loss, wake_phi_loss)
        if isinstance(loss, tuple):
            return torch_item(loss[0]), torch_item(loss[1])
        else:
            return torch_item(loss)

1	import warnings	4✔
2
3	import torch	4✔
4
5	import pyro	4✔
6	import pyro.optim	4✔
7	import pyro.poutine as poutine	4✔
8	from pyro.infer.abstract_infer import TracePosterior	4✔
9	from pyro.infer.elbo import ELBO	4✔
10	from pyro.infer.util import torch_item	4✔
11
12
13	class SVI(TracePosterior):	4✔
14	"""
15	:param model: the model (callable containing Pyro primitives)
16	:param guide: the guide (callable containing Pyro primitives)
17	:param optim: a wrapper a for a PyTorch optimizer
18	:type optim: pyro.optim.PyroOptim
19	:param loss: an instance of a subclass of :class:`~pyro.infer.elbo.ELBO`.
20	Pyro provides three built-in losses:
21	:class:`~pyro.infer.trace_elbo.Trace_ELBO`,
22	:class:`~pyro.infer.tracegraph_elbo.TraceGraph_ELBO`, and
23	:class:`~pyro.infer.traceenum_elbo.TraceEnum_ELBO`.
24	See the :class:`~pyro.infer.elbo.ELBO` docs to learn how to implement
25	a custom loss.
26	:type loss: pyro.infer.elbo.ELBO
27	:param num_samples: (DEPRECATED) the number of samples for Monte Carlo posterior approximation
28	:param num_steps: (DEPRECATED) the number of optimization steps to take in ``run()``
29
30	A unified interface for stochastic variational inference in Pyro. The most
31	commonly used loss is ``loss=Trace_ELBO()``. See the tutorial
32	`SVI Part I <http://pyro.ai/examples/svi_part_i.html>`_ for a discussion.
33	"""
34	def __init__(self,	4✔
35	model,
36	guide,
37	optim,
38	loss,
39	loss_and_grads=None,
40	num_samples=0,
41	num_steps=0,
42	**kwargs):
43	if num_steps:	4✔
44	warnings.warn('The `num_steps` argument to SVI is deprecated and will be removed in '
45	'a future release. Use `SVI.step` directly to control the '
46	'number of iterations.', FutureWarning)
47	if num_samples:	4✔
48	warnings.warn('The `num_samples` argument to SVI is deprecated and will be removed in '
49	'a future release. Use `pyro.infer.Predictive` class to draw '
50	'samples from the posterior.', FutureWarning)
51
52	self.model = model	4✔
53	self.guide = guide	4✔
54	self.optim = optim	4✔
55	self.num_steps = num_steps	4✔
56	self.num_samples = num_samples	4✔
57	super(SVI, self).__init__(**kwargs)	4✔
58
59	if not isinstance(optim, pyro.optim.PyroOptim):	4✔
60	raise ValueError("Optimizer should be an instance of pyro.optim.PyroOptim class.")
61
62	if isinstance(loss, ELBO):	4✔
63	self.loss = loss.loss	4✔
64	self.loss_and_grads = loss.loss_and_grads	4✔
65	else:
66	if loss_and_grads is None:	3✔
67	def _loss_and_grads(args, *kwargs):	2✔
68	loss_val = loss(args, *kwargs)	2✔
69	loss_val.backward(retain_graph=True)	2✔
70	return loss_val	2✔
71	loss_and_grads = _loss_and_grads	2✔
72	self.loss = loss	3✔
73	self.loss_and_grads = loss_and_grads	3✔
74
75	def run(self, args, *kwargs):	4✔
76	"""
77	.. warning::
78	This method is deprecated, and will be removed in a future release.
79	For inference, use :meth:`step` directly, and for predictions,
80	use the :class:`~pyro.infer.predictive.Predictive` class.
81	"""
82	warnings.warn('The `SVI.run` method is deprecated and will be removed in a '
83	'future release. For inference, use `SVI.step` directly, '
84	'and for predictions, use the `pyro.infer.Predictive` class.',
85	FutureWarning)
86	if self.num_steps > 0:	1✔
87	with poutine.block():	×
88	for i in range(self.num_steps):	×
89	self.step(args, *kwargs)	×
90	return super(SVI, self).run(args, *kwargs)	1✔
91
92	def _traces(self, args, *kwargs):	4✔
93	for i in range(self.num_samples):	1✔
94	guide_trace = poutine.trace(self.guide).get_trace(args, *kwargs)	1✔
95	model_trace = poutine.trace(poutine.replay(self.model, trace=guide_trace)).get_trace(args, *kwargs)	1✔
96	yield model_trace, 1.0	1✔
97
98	def evaluate_loss(self, args, *kwargs):	4✔
99	"""
100	:returns: estimate of the loss
101	:rtype: float
102
103	Evaluate the loss function. Any args or kwargs are passed to the model and guide.
104	"""
105	with torch.no_grad():	1✔
106	loss = self.loss(self.model, self.guide, args, *kwargs)	1✔
107	if isinstance(loss, tuple):	1✔
108	return torch_item(loss[0]), torch_item(loss[1])	×
109	else:
110	return torch_item(loss)	1✔
111
112	def step(self, args, *kwargs):	4✔
113	"""
114	:returns: estimate of the loss
115	:rtype: float
116
117	Take a gradient step on the loss function (and any auxiliary loss functions
118	generated under the hood by `loss_and_grads`).
119	Any args or kwargs are passed to the model and guide
120	"""
121	# get loss and compute gradients
122	with poutine.trace(param_only=True) as param_capture:	4✔
123	loss = self.loss_and_grads(self.model, self.guide, args, *kwargs)	4✔
124
125	params = set(site["value"].unconstrained()	4✔
126	for site in param_capture.trace.nodes.values())
127
128	# actually perform gradient steps
129	# torch.optim objects gets instantiated for any params that haven't been seen yet
130	self.optim(params)	4✔
131
132	# zero gradients
133	pyro.infer.util.zero_grads(params)	4✔
134
135	# RWS: handle loss = (wake_theta_loss, wake_phi_loss)
136	if isinstance(loss, tuple):	4✔
137	return torch_item(loss[0]), torch_item(loss[1])	2✔
138	else:
139	return torch_item(loss)	4✔

pyro-ppl / pyro / 5561

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous