Add clip_grads_by_norm (#312)

luchador/agent/rl/q_learning.py

@@ -45,36 +45,13 @@ def _build_error(target_q, action_value_0, action):
     return nn.ops.reduce_sum(mask * error, axis=1)
 
 
-def _clip_grads(grads_and_vars, clip_norm):
-    """Clip gradients by norm
-
-    Parameters
-    ----------
-    grads_and_vars : list
-        Gradient and Variable tuples. Return value from ``compute_gradients``.
-
-    clip_norm : Number or Tensor
-        Value to clip gradients
-
-    Returns
-    -------
-    list
-        Resulting gradients and vars pairs
-    """
-    ret = []
-    for grad, var in grads_and_vars:
-        name = '{}_clip'.format(grad.name)
-        grad = nn.ops.clip_by_norm(grad, clip_norm=clip_norm, name=name)
-        ret.append((grad, var))
-    return ret
-
-
 def _build_optimize_op(optimizer, loss, params, clip_grads):
     grads_and_vars = nn.ops.compute_gradient(loss=loss, wrt=params)
     # Remove untrainable variables
     grads_and_vars = [g_v for g_v in grads_and_vars if g_v[0] is not None]
     if clip_grads:
-        grads_and_vars = _clip_grads(grads_and_vars, **clip_grads)
+        grads_and_vars = nn.ops.clip_grads_by_norm(
+            grads_and_vars, **clip_grads)
     return optimizer.apply_gradients(grads_and_vars)
 
 

luchador/nn/core/impl/ops/__init__.py

@@ -1,7 +1,7 @@
 """Define operations"""
 from __future__ import absolute_import
 # pylint: disable=unused-import,redefined-builtin
-from .clip import clip_by_value, clip_by_norm  # noqa
+from .clip import clip_by_value, clip_by_norm, clip_grads_by_norm  # noqa
 from .grad import compute_gradient  # noqa
 from .misc import build_sync_op, one_hot  # noqa
 from .transform import reshape, tile  # noqa
@@ -13,7 +13,7 @@
 )
 
 __all__ = [
-    'clip_by_value', 'clip_by_norm',
+    'clip_by_value', 'clip_by_norm', 'clip_grads_by_norm',
     'add', 'multiply', 'maximum', 'minimum',
     'compute_gradient',
     'dot',

luchador/nn/core/impl/ops/clip.py

@@ -4,7 +4,7 @@
 from luchador.nn.core.base import BaseWrapper
 from ... import backend as be
 
-__all__ = ['clip_by_value', 'clip_by_norm']
+__all__ = ['clip_by_value', 'clip_by_norm', 'clip_grads_by_norm']
 
 
 def _is_wrapper(obj):
@@ -61,3 +61,28 @@ def clip_by_norm(tensor, clip_norm, axes=None, name=None):
     if _is_wrapper(clip_norm):
         clip_norm = clip_norm.unwrap()
     return be.ops.clip_by_norm(tensor, clip_norm, axes, name)
+
+
+def clip_grads_by_norm(grads_and_vars, clip_norm):
+    """Clip each gradient by norm
+
+    Parameters
+    ----------
+    grads_and_vars : list
+        Gradient and Variable tuples. Return value from
+        :func:`luchador.nn.ops.compute_gradients`.
+
+    clip_norm : Number or Tensor
+        Value to clip gradients
+
+    Returns
+    -------
+    list
+        Resulting gradients and vars pairs
+    """
+    ret = []
+    for grad, var in grads_and_vars:
+        name = '{}_clip'.format(grad.name)
+        grad = clip_by_norm(grad, clip_norm=clip_norm, name=name)
+        ret.append((grad, var))
+    return ret