Merge pull request #122 from Yoctol/clip-global-norm

Clip global norm

talos/__version__.py

@@ -1,4 +1,4 @@
 __title__ = 'talos'
-__version__ = '1.6.3'
+__version__ = '1.6.4'
 __description__ = 'Powerful Neural Network Builder'
 __author__ = 'Jsaon'

talos/optimizers/gradient_clipping.py

@@ -3,7 +3,7 @@
 
 class GradientClipping(tf.train.Optimizer):
 
-    _ALLOWED_CLIP_BY = {'value', 'norm'}
+    _ALLOWED_CLIP_BY = {'value', 'norm', 'global_norm'}
 
     def __init__(
             self,
@@ -25,10 +25,21 @@ def __init__(
         self.clip_by = clip_by
 
     def apply_gradients(self, grads_and_vars, global_step=None, name=None):
-        processed_gvs = [
-            (self._process_grad(g), v) for g, v in grads_and_vars
-            if g is not None
-        ]
+        if self.clip_by == 'global_norm':
+            gs, vs = [], []
+            for g, v in grads_and_vars:
+                if g is not None:
+                    gs.append(g)
+                    vs.append(v)
+
+            clipped_g, _ = tf.clip_by_global_norm(gs, clip_norm=self.value)
+            processed_gvs = list(zip(clipped_g, vs))
+        else:
+            processed_gvs = [
+                (self._process_grad(g), v) for g, v in grads_and_vars
+                if g is not None
+            ]
+
         return self.optimizer.apply_gradients(
             processed_gvs,
             global_step=global_step,

talos/optimizers/tests/test_gradient_clipping.py

@@ -42,3 +42,31 @@ def test_clip_norm(sess):
         sess.run(x),
         x_val - lr * x_val * np.minimum(value / np.linalg.norm(x_val), 1.),
     )
+
+
+def test_clip_global_norm(sess):
+    lr, value = 0.2, 0.5
+    x_val, y_val = np.array([3., 4.]), np.array([[1.], [2.]])
+    optimizer = GradientClipping(
+        tf.train.GradientDescentOptimizer(lr),
+        value,
+        clip_by='global_norm',
+    )
+    x = tf.Variable(x_val)
+    y = tf.Variable(y_val)
+    z = tf.nn.l2_loss(x) + tf.nn.l2_loss(y)  # dz/dx = x, dz/dy = y
+
+    train_op = optimizer.minimize(z)
+    sess.run(tf.variables_initializer([x, y]))
+    sess.run(train_op)
+
+    new_x_val, new_y_val = sess.run([x, y])
+    global_norm = np.sqrt(np.sum(x_val ** 2) + np.sum(y_val ** 2))
+    np.testing.assert_array_almost_equal(
+        new_x_val,
+        x_val - lr * x_val * np.minimum(value / global_norm, 1.),
+    )
+    np.testing.assert_array_almost_equal(
+        new_y_val,
+        y_val - lr * y_val * np.minimum(value / global_norm, 1.),
+    )