Optimization

pymc3/__init__.py

@@ -27,6 +27,7 @@
 from .tests import test
 
 from .data import *
+from .optimization import *
 
 import logging
 _log = logging.getLogger('pymc3')

pymc3/optimization.py

@@ -0,0 +1,83 @@
+import tqdm
+import numpy as np
+import theano
+from theano.configparser import change_flags
+import pymc3 as pm
+
+__all__ = [
+    'Optimizer'
+]
+
+
+class Optimizer(object):
+    """
+    Optimization with posterior replacements
+
+    Parameters
+    ----------
+    approx : Approximation
+    loss : scalar
+    wrt : shared params
+    more_replacements : other replacements in the graph
+    optimizer : callable that returns updates, pm.adam by default
+    """
+    @change_flags(compute_test_value='off')
+    def __init__(self, approx, loss, wrt, more_replacements=None, optimizer=None):
+        if optimizer is None:
+            optimizer = pm.adam
+        self.optimizer = optimizer
+        self.wrt = wrt
+        self.approx = approx
+        self.loss = approx.apply_replacements(loss, more_replacements=more_replacements)
+        updates = self.optimizer(self.loss, self.wrt)
+        self.step_function = theano.function([], self.loss, updates=updates)
+        self.hist = np.asarray(())
+
+    @change_flags(compute_test_value='off')
+    def refresh(self, kwargs=None):
+        """
+        Recompile step_function and reset updates
+
+        This can be needed sometimes when
+        you use stateful optimizers like ADAM
+        and want to reuse the function for new problem
+
+        Parameters
+        ----------
+        kwargs : kwargs for theano.function
+        """
+        updates = self.optimizer(self.loss, self.wrt)
+        self.step_function = theano.function([], self.loss, updates=updates, **kwargs)
+        self.hist = np.asarray(())
+
+    def fit(self, n=5000, callbacks=()):
+        """
+        Perform optimization steps
+
+        Parameters
+        ----------
+        n : int
+            number of iterations
+        callbacks : list[callable]
+            list of callables with following signature
+            f(Approximation, loss_history, i) -> None
+        """
+        progress = tqdm.trange(n)
+        scores = np.empty(n)
+        scores[:] = np.nan
+        i = 0
+        try:
+            for i in progress:
+                scores[i] = self.step_function()
+                for callback in callbacks:
+                    callback(self.approx, scores[:i + 1], i)
+                if i % ((n+1000)//1000) == 0:
+                    progress.set_description(
+                        'E_q[Loss] = %.4f' % scores[max(0, i - n//50):i+1].mean()
+                    )
+        except (KeyboardInterrupt, StopIteration):
+            self.hist = np.concatenate((self.hist, scores[:i]))
+        else:
+            self.hist = np.concatenate((self.hist, scores))
+        finally:
+            progress.close()

pymc3/variational/opvi.py

@@ -593,11 +593,11 @@ def apply_replacements(self, node, deterministic=False,
 
         Parameters
         ----------
-        node : Theano Variables (or Theano expressions)
+        node : Theano Variable(s)
             node or nodes for replacements
         deterministic : bool
-            whether to use zeros as initial distribution
-            if True - zero initial point will produce constant latent variables
+            whether to use point with highest density as initial point for distribution
+            if True - constant initial point will produce constant latent variables
         include : list
             latent variables to be replaced
         exclude : list
@@ -623,7 +623,7 @@ def sample_node(self, node, size=100,
 
         Parameters
         ----------
-        node : Theano Variables (or Theano expressions)
+        node : Theano Variable(s)
         size : scalar
             number of samples
         more_replacements : dict