add class based optimizers

Author: ferrine

pymc3/variational/updates.py

@@ -97,7 +97,7 @@
 """
 
 from collections import OrderedDict
-
+import functools
 import numpy as np
 
 import theano
@@ -117,10 +117,28 @@
     "adam",
     "adamax",
     "norm_constraint",
-    "total_norm_constraint"
+    "total_norm_constraint",
+    "Sgd",
+    "Momentum",
+    "NesterovMomentum",
+    "Adagrad",
+    "RMSProp",
+    "AdaDelta",
+    "AdaMax",
+    "Adam",
 ]
 
 
+class Optimizer(object):
+    _opt = None
+
+    def __init__(self, *args, **kwargs):
+        self.opt = functools.partial(self._opt, *args, **kwargs)
+
+    def __call__(self, loss_or_grads, params):
+        return self.opt(loss_or_grads, params)
+
+
 def get_or_compute_grads(loss_or_grads, params):
     """Helper function returning a list of gradients
 
@@ -160,7 +178,7 @@ def get_or_compute_grads(loss_or_grads, params):
         return theano.grad(loss_or_grads, params)
 
 
-def sgd(loss_or_grads, params, learning_rate):
+def sgd(loss_or_grads, params, learning_rate=1e-3):
     """Stochastic Gradient Descent (SGD) updates
 
     Generates update expressions of the form:
@@ -190,6 +208,10 @@ def sgd(loss_or_grads, params, learning_rate):
     return updates
 
 
+class Sgd(Optimizer):
+    _opt = sgd
+
+
 def apply_momentum(updates, params=None, momentum=0.9):
     """Returns a modified update dictionary including momentum
 
@@ -277,6 +299,10 @@ def momentum(loss_or_grads, params, learning_rate, momentum=0.9):
     return apply_momentum(updates, momentum=momentum)
 
 
+class Momentum(Optimizer):
+    _opt = momentum
+
+
 def apply_nesterov_momentum(updates, params=None, momentum=0.9):
     """Returns a modified update dictionary including Nesterov momentum
 
@@ -331,7 +357,7 @@ def apply_nesterov_momentum(updates, params=None, momentum=0.9):
     return updates
 
 
-def nesterov_momentum(loss_or_grads, params, learning_rate, momentum=0.9):
+def nesterov_momentum(loss_or_grads, params, learning_rate=1e-3, momentum=0.9):
     """Stochastic Gradient Descent (SGD) updates with Nesterov momentum
 
     Generates update expressions of the form:
@@ -375,6 +401,10 @@ def nesterov_momentum(loss_or_grads, params, learning_rate, momentum=0.9):
     return apply_nesterov_momentum(updates, momentum=momentum)
 
 
+class NesterovMomentum(Optimizer):
+    _opt = nesterov_momentum
+
+
 def adagrad(loss_or_grads, params, learning_rate=1.0, epsilon=1e-6):
     """Adagrad updates
 
@@ -434,6 +464,10 @@ def adagrad(loss_or_grads, params, learning_rate=1.0, epsilon=1e-6):
     return updates
 
 
+class Adagrad(Optimizer):
+    _opt = adagrad
+
+
 def rmsprop(loss_or_grads, params, learning_rate=1.0, rho=0.9, epsilon=1e-6):
     """RMSProp updates
 
@@ -495,6 +529,10 @@ def rmsprop(loss_or_grads, params, learning_rate=1.0, rho=0.9, epsilon=1e-6):
     return updates
 
 
+class RMSProp(Optimizer):
+    _opt = rmsprop
+
+
 def adadelta(loss_or_grads, params, learning_rate=1.0, rho=0.95, epsilon=1e-6):
     """ Adadelta updates
 
@@ -579,6 +617,10 @@ def adadelta(loss_or_grads, params, learning_rate=1.0, rho=0.95, epsilon=1e-6):
     return updates
 
 
+class AdaDelta(Optimizer):
+    _opt = adadelta
+
+
 def adam(loss_or_grads, params, learning_rate=0.001, beta1=0.9,
          beta2=0.999, epsilon=1e-8):
     """Adam updates
@@ -646,6 +688,10 @@ def adam(loss_or_grads, params, learning_rate=0.001, beta1=0.9,
     return updates
 
 
+class Adam(Optimizer):
+    _opt = adam
+
+
 def adamax(loss_or_grads, params, learning_rate=0.002, beta1=0.9,
            beta2=0.999, epsilon=1e-8):
     """Adamax updates
@@ -708,6 +754,10 @@ def adamax(loss_or_grads, params, learning_rate=0.002, beta1=0.9,
     return updates
 
 
+class AdaMax(Optimizer):
+    _opt = adamax
+
+
 def norm_constraint(tensor_var, max_norm, norm_axes=None, epsilon=1e-7):
     """Max weight norm constraints and gradient clipping