Add LearningRateMultiplier wrapper for optimizers

keras_contrib/optimizers/__init__.py

@@ -1,6 +1,7 @@
 from .ftml import FTML
 from .padam import Padam
 from .yogi import Yogi
+from .lr_multiplier import LearningRateMultiplier
 
 # aliases
 ftml = FTML

keras_contrib/optimizers/lr_multiplier.py

@@ -0,0 +1,66 @@
+from keras.optimizers import Optimizer
+from keras.utils import get_custom_objects
+
+
+class LearningRateMultiplier(Optimizer):
+    """Optimizer wrapper for per layer learning rate.
+
+    This wrapper is used to add per layer learning rates by
+    providing per layer factors which are multiplied with the
+    learning rate of the optimizer.
+
+    Note: This is a wrapper and does not implement any
+    optimization algorithm.
+
+    # Arguments
+        optimizer: An optimizer class to be wrapped.
+        lr_multipliers: Dictionary of the per layer factors. For
+            example `optimizer={'conv_1/kernel':0.5, 'conv_1/bias':0.1}`.
+            If for kernel and bias the same learning rate is used, the
+            user can specify `optimizer={'conv_1':0.5}`.
+        **kwargs: The arguments for instantiating the wrapped optimizer
+            class.
+    """
+    def __init__(self, optimizer, lr_multipliers=None, **kwargs):
+        self._class = optimizer
+        self._optimizer = optimizer(**kwargs)
+        self._lr_multipliers = lr_multipliers or {}
+
+    def _get_multiplier(self, param):
+        for k in self._lr_multipliers.keys():
+            if k in param.name:
+                return self._lr_multipliers[k]
+
+    def get_updates(self, loss, params):
+        mult_lr_params = {p: self._get_multiplier(p) for p in params
+                          if self._get_multiplier(p)}
+        base_lr_params = [p for p in params if self._get_multiplier(p) is None]
+
+        updates = []
+        base_lr = self._optimizer.lr
+        for param, multiplier in mult_lr_params.items():
+            self._optimizer.lr = base_lr * multiplier
+            updates.extend(self._optimizer.get_updates(loss, [param]))
+
+        self._optimizer.lr = base_lr
+        updates.extend(self._optimizer.get_updates(loss, base_lr_params))
+
+        return updates
+
+    def get_config(self):
+        config = {'optimizer': self._class,
+                  'lr_multipliers': self._lr_multipliers}
+        base_config = self._optimizer.get_config()
+        return dict(list(base_config.items()) + list(config.items()))
+
+    def __getattr__(self, name):
+        return getattr(self._optimizer, name)
+
+    def __setattr__(self, name, value):
+        if name.startswith('_'):
+            super(LearningRateMultiplier, self).__setattr__(name, value)
+        else:
+            self._optimizer.__setattr__(name, value)
+
+
+get_custom_objects().update({'LearningRateMultiplier': LearningRateMultiplier})

tests/keras_contrib/optimizers/lr_multiplier_test.py

@@ -0,0 +1,14 @@
+from keras_contrib.tests import optimizers
+from keras_contrib.optimizers import LearningRateMultiplier
+from keras.optimizers import SGD, Adam
+from keras.callbacks import LearningRateScheduler
+
+
+def test_lr_multiplier():
+    opt1 = LearningRateMultiplier(SGD, lr=0.01, momentum=0.9, nesterov=True)
+    optimizers._test_optimizer(opt1, target=0.95)
+
+    mult = {'dense': 10}
+    opt2 = LearningRateMultiplier(SGD, lr_multipliers=mult,
+                                  lr=0.001, momentum=0.9, nesterov=True)
+    optimizers._test_optimizer(opt2, target=0.95)