Remove names from register_hook calls (pytorch#446)

colesbury · soumith · commit 7e4ddcfe8aad · 2017-01-13T15:57:03.000-05:00
The register hook calls now return an object that can be used to remove
the hook. For example,

   &gt;&gt;&gt; h = module.register_forward_hook(callback)
   &gt;&gt;&gt; h.remove()  # removes hook

Or as a context manager:

   &gt;&gt;&gt; with module.register_forward_hook(callback):
   ...     pass

This makes it easier for libraries to use hooks without worrying about
name collisions.
diff --git a/test/test_autograd.py b/test/test_autograd.py
@@ -72,28 +72,28 @@ def bw_hook(inc, grad):
             counter[0] += inc
 
         z = x ** 2 + x * 2 + x * y + y
-        z.register_hook('test', lambda *args: bw_hook(1, *args))
+        test = z.register_hook(lambda *args: bw_hook(1, *args))
         z.backward(torch.ones(5, 5), retain_variables=True)
         self.assertEqual(counter[0], 1)
 
-        z.register_hook('test2', lambda *args: bw_hook(2, *args))
+        test2 = z.register_hook(lambda *args: bw_hook(2, *args))
         z.backward(torch.ones(5, 5), retain_variables=True)
         self.assertEqual(counter[0], 4)
 
-        z.remove_hook('test2')
+        test2.remove()
         z.backward(torch.ones(5, 5), retain_variables=True)
         self.assertEqual(counter[0], 5)
 
         def bw_hook_modify(grad):
             return grad.mul(2)
 
-        z.remove_hook('test')
-        z.register_hook('test', bw_hook_modify)
+        test.remove()
+        z.register_hook(bw_hook_modify)
         y.grad.zero_()
         z.backward(torch.ones(5, 5), retain_variables=True)
         self.assertEqual(y.grad, (x.data + 1) * 2)
 
-        y.register_hook('test', bw_hook_modify)
+        y.register_hook(bw_hook_modify)
         y.grad.zero_()
         z.backward(torch.ones(5, 5))
         self.assertEqual(y.grad, (x.data + 1) * 4)
@@ -420,8 +420,7 @@ def backward(self, grad_a, grad_b):
 
         q, p = Identity()(x, y)
         # Make sure hooks only receive grad from usage of q, not x.
-        q.register_hook(
-                'test', lambda grad: self.assertEqual(grad, torch.ones(5, 5)))
+        q.register_hook(lambda grad: self.assertEqual(grad, torch.ones(5, 5)))
         (q + p + x).sum().backward()
         self.assertEqual(x.grad, torch.ones(5, 5) * 3)
         self.assertEqual(y.grad, torch.ones(5, 5))
diff --git a/test/test_nn.py b/test/test_nn.py
@@ -227,14 +227,14 @@ def bw_hook(inc, h_module, grad_input, grad_output):
             self.assertEqual(grad_output[0], torch.ones(5, 5) * 2)
             counter['backwards'] += inc
 
-        module.register_forward_hook('test', lambda *args: fw_hook(1, *args))
+        test_fwd = module.register_forward_hook(lambda *args: fw_hook(1, *args))
 
         module(input)
         module(input)
         self.assertEqual(counter['forwards'], 2)
         self.assertEqual(counter['backwards'], 0)
 
-        module.register_backward_hook('test', lambda *args: bw_hook(1, *args))
+        test_bwd = module.register_backward_hook(lambda *args: bw_hook(1, *args))
 
         output = module(input)
         self.assertEqual(counter['forwards'], 3)
@@ -248,32 +248,32 @@ def bw_hook(inc, h_module, grad_input, grad_output):
         self.assertEqual(counter['forwards'], 3)
         self.assertEqual(counter['backwards'], 2)
 
-        module.register_forward_hook('test2', lambda *args: fw_hook(2, *args))
+        test2_fwd = module.register_forward_hook(lambda *args: fw_hook(2, *args))
 
         output = module(input)
         self.assertEqual(counter['forwards'], 6)
         self.assertEqual(counter['backwards'], 2)
 
-        module.register_backward_hook('test2', lambda *args: bw_hook(2, *args))
+        test2_bwd = module.register_backward_hook(lambda *args: bw_hook(2, *args))
 
         module(input).backward(torch.ones(5, 5) * 2)
         self.assertEqual(counter['forwards'], 9)
         self.assertEqual(counter['backwards'], 5)
 
-        module.remove_backward_hook('test2')
+        test2_bwd.remove()
 
         module(input).backward(torch.ones(5, 5) * 2)
         self.assertEqual(counter['forwards'], 12)
         self.assertEqual(counter['backwards'], 6)
 
-        module.remove_forward_hook('test2')
+        test2_fwd.remove()
 
         module(input).backward(torch.ones(5, 5) * 2)
         self.assertEqual(counter['forwards'], 13)
         self.assertEqual(counter['backwards'], 7)
 
-        module.remove_forward_hook('test')
-        module.remove_backward_hook('test')
+        test_fwd.remove()
+        test_bwd.remove()
 
     def test_hook_fail(self):
         module = nn.Sigmoid()
@@ -291,33 +291,29 @@ def bw_fail1(self, grad_input, grad_output):
         def bw_fail2(self, grad_input, grad_output):
             return grad_input + (torch.randn(2, 2),)
 
-        module.register_forward_hook('fw_fail', fw_fail1)
-        with self.assertRaises(RuntimeError) as err:
-            module(input)
-        self.assertIn("fw_fail", err.exception.args[0])
-        self.assertIn("didn't return None", err.exception.args[0])
-        module.remove_forward_hook('fw_fail')
+        with module.register_forward_hook(fw_fail1):
+            with self.assertRaises(RuntimeError) as err:
+                module(input)
+            self.assertIn("fw_fail", err.exception.args[0])
+            self.assertIn("didn't return None", err.exception.args[0])
 
-        module.register_forward_hook('fw_fail2', fw_fail2)
-        with self.assertRaises(RuntimeError) as err:
-            module(input)
-        self.assertIn("fw_fail2", err.exception.args[0])
-        self.assertIn("didn't return None", err.exception.args[0])
-        module.remove_forward_hook('fw_fail2')
-
-        module.register_backward_hook('bw_fail', bw_fail1)
-        with self.assertRaises(RuntimeError) as err:
-            module(input).sum().backward()
-        self.assertIn("bw_fail", err.exception.args[0])
-        self.assertIn("got 0, but expected 1", err.exception.args[0])
-        module.remove_backward_hook('bw_fail')
-
-        module.register_backward_hook('bw_fail2', bw_fail2)
-        with self.assertRaises(RuntimeError) as err:
-            module(input).sum().backward()
-        self.assertIn("bw_fail2", err.exception.args[0])
-        self.assertIn("got 2, but expected 1", err.exception.args[0])
-        module.remove_backward_hook('bw_fail2')
+        with module.register_forward_hook(fw_fail2):
+            with self.assertRaises(RuntimeError) as err:
+                module(input)
+            self.assertIn("fw_fail2", err.exception.args[0])
+            self.assertIn("didn't return None", err.exception.args[0])
+
+        with module.register_backward_hook(bw_fail1):
+            with self.assertRaises(RuntimeError) as err:
+                module(input).sum().backward()
+            self.assertIn("bw_fail", err.exception.args[0])
+            self.assertIn("got 0, but expected 1", err.exception.args[0])
+
+        with module.register_backward_hook(bw_fail2):
+            with self.assertRaises(RuntimeError) as err:
+                module(input).sum().backward()
+            self.assertIn("bw_fail2", err.exception.args[0])
+            self.assertIn("got 2, but expected 1", err.exception.args[0])
 
     def test_hook_writeable(self):
         module = nn.Linear(5, 5)
@@ -326,7 +322,7 @@ def test_hook_writeable(self):
         def bw_hook(self, grad_input, grad_output):
             return tuple(gi * 2 for gi in grad_input)
 
-        module.register_backward_hook('test', bw_hook)
+        module.register_backward_hook(bw_hook)
         module(input).backward(torch.ones(5, 5))
         expected_grad = torch.ones(5, 5).mm(module.weight.data) * 2
         self.assertEqual(input.grad, expected_grad)
diff --git a/torch/autograd/function.py b/torch/autograd/function.py
@@ -1,5 +1,6 @@
 import torch
 import torch._C as _C
+import torch.utils.hooks as hooks
 from collections import OrderedDict
 from itertools import chain
 
@@ -106,16 +107,12 @@ def mark_non_differentiable(self, *args):
         """
         self.non_differentiable = args
 
-    def register_hook(self, name, hook):
-        self._backward_hooks = self._backward_hooks or OrderedDict()
-        assert name not in self._backward_hooks, \
-            "Trying to register a second hook with name {}".format(name)
-        self._backward_hooks[name] = hook
-
-    def remove_hook(self, name):
-        assert self._backward_hooks and name in self._backward_hooks, \
-            "Trying to remove an inexistent hook with name {}".format(name)
-        del self._backward_hooks[name]
+    def register_hook(self, hook):
+        if self._backward_hooks is None:
+            self._backward_hooks = OrderedDict()
+        handle = hooks.RemovableHandle(self._backward_hooks)
+        self._backward_hooks[id(handle)] = hook
+        return handle
 
     def forward(self, *input):
         """Performs the operation.
diff --git a/torch/autograd/variable.py b/torch/autograd/variable.py
@@ -1,6 +1,7 @@
 import sys
 import torch._C as _C
 from collections import OrderedDict
+import torch.utils.hooks as hooks
 
 from ._functions import *
 
@@ -155,54 +156,43 @@ def backward(self, gradient=None, retain_variables=False):
             gradient = self.data.new().resize_as_(self.data).fill_(1)
         self._execution_engine.run_backward((self,), (gradient,), retain_variables)
 
-    def register_hook(self, name, hook):
-        """Registers a named backward hook.
+    def register_hook(self, hook):
+        """Registers a backward hook.
 
-        Given hook will be saved and called with the gradient w.r.t. the
-        variable at every backward pass. To remove a hook use
-        :func:`remove_hook`. Saved hooks are called in the same order
-        in which they were registered.
+        The hook will be called every time a gradient with respect to the
+        variable is computed. The hook should have the following signature::
 
-        You should never modify the data of gradient tensor given to your hook,
-        but you can use it in out-of-place operations and return a new tensor::
+            hook(grad) -> Tensor or None
 
-            variable.register_hook('double_grad', lambda grad: grad * 2)
+        The hook should not modify its argument, but it can optionally return
+        a new gradient which will be used in place of :attr:`grad`.
 
-        The returned value will replace the original tensor. Note that you
-        don't need to return anything from the hook, in which case it won't
-        change the gradient.
+        This function returns a handle with a method ``handle.remove()``
+        that removes the hook from the module.
 
-
-        Parameters:
-            name(str): Name of the hook.
-            hook(callable): Hook callable. It will be given a single argument
-                that's gradient w.r.t. the variable its registered on.
+        Example:
+            >>> v = Variable(torch.Tensor([0, 0, 0]), requires_grad=True)
+            >>> h = v.register_hook(lambda grad: grad * 2)  # double the gradient
+            >>> v.backward(torch.Tensor([1, 1, 1]))
+            >>> v.grad
+             2
+             2
+             2
+            [torch.FloatTensor of size 3]
+            >>> h.remove()  # removes the hook
         """
         if self.volatile:
-            raise RuntimeError("registering hook on a volatile variable")
+            raise RuntimeError("cannot register a hook on a volatile variable")
         if not self.requires_grad:
-            raise RuntimeError("registering hook on a variable that doesn't require gradient")
+            raise RuntimeError("cannot register a hook on a variable that "
+                               "doesn't require gradient")
         if self._backward_hooks is None:
             self._backward_hooks = OrderedDict()
             if self.creator is not None:
                 self.creator._register_hook_dict(self)
-        assert name not in self._backward_hooks, \
-            "Trying to register a second hook with name {}".format(name)
-        self._backward_hooks[name] = hook
-
-    def remove_hook(self, name):
-        """Removes a previously registered backward hook.
-
-        Raises RuntimeError if there's no hook registered under a given name.
-
-        Parameters:
-            name(str): Name of the hook.
-        """
-        if self.volatile:
-            raise RuntimeError("volatile variables don't support hooks")
-        assert self._backward_hooks and name in self._backward_hooks, \
-            "Trying to remove an inexistent hook with name {}".format(name)
-        del self._backward_hooks[name]
+        handle = hooks.RemovableHandle(self._backward_hooks)
+        self._backward_hooks[id(handle)] = hook
+        return handle
 
     def _do_backward(self, grad_output, retain_variables):
         assert len(grad_output) == 1
diff --git a/torch/csrc/autograd/function.cpp b/torch/csrc/autograd/function.cpp
@@ -461,14 +461,15 @@ PyObject *THPFunction_do_forward(THPFunction *self, PyObject *inputs)
 
 // We need a reference to a smart pointer that will outlive the duration of
 // a function call, so that the char* pointer is valid even after it returns
-static char* _try_get_name(PyObject *key, THPObjectPtr& tmp) {
+static char* _try_get_name(PyObject *hook, THPObjectPtr& tmp) {
+  tmp = PyObject_GetAttrString(hook, "__name__");
 #if PY_MAJOR_VERSION == 2
-  if (PyString_Check(key)) {
-    return PyString_AS_STRING(key);
+  if (tmp && PyString_Check(tmp.get())) {
+    return PyString_AS_STRING(tmp.get());
   }
 #else
-  if (PyUnicode_Check(key)) {
-    tmp = PyUnicode_AsASCIIString(key);
+  if (tmp && PyUnicode_Check(tmp.get())) {
+    tmp = PyUnicode_AsASCIIString(tmp.get());
     return PyBytes_AS_STRING(tmp.get());
   }
 #endif
@@ -481,7 +482,7 @@ static char* _try_get_name(PyObject *key, THPObjectPtr& tmp) {
   hook_name ? "' " : ""
 
 static void _ensure_correct_hook_result_single(PyObject *original,
-    PyObject *returned, PyObject *key)
+    PyObject *returned, PyObject *hook)
 {
 #if PY_MAJOR_VERSION == 2
   static PyObject *IS_SAME_SIZE_NAME = PyString_FromString("is_same_size");
@@ -491,7 +492,7 @@ static void _ensure_correct_hook_result_single(PyObject *original,
   THPObjectPtr tmp;
   // Check that the type matches
   if(Py_TYPE(original) != Py_TYPE(returned)) {
-    char *hook_name = _try_get_name(key, tmp);
+    char *hook_name = _try_get_name(hook, tmp);
     THPUtils_setError("backward hook %s%s%shas changed the type of "
         "grad_input (was %s, but got %s)",
         OPTIONAL_HOOK_NAME,
@@ -505,7 +506,7 @@ static void _ensure_correct_hook_result_single(PyObject *original,
   THPObjectPtr is_same_size = PyObject_CallMethodObjArgs(original,
       IS_SAME_SIZE_NAME, returned, NULL);
   if(is_same_size.get() != Py_True) {
-    char *hook_name = _try_get_name(key, tmp);
+    char *hook_name = _try_get_name(hook, tmp);
     THPUtils_setError("backward hook %s%s%shas changed the size of "
         "grad_input",
         OPTIONAL_HOOK_NAME
@@ -515,12 +516,12 @@ static void _ensure_correct_hook_result_single(PyObject *original,
 }
 
 static void _ensure_correct_hook_result(THPObjectPtr& grad_input,
-    THPObjectPtr& result, PyObject *key)
+    THPObjectPtr& result, PyObject *hook)
 {
   THPObjectPtr tmp;
   // Check that the tuple sizes match
   if (PyTuple_GET_SIZE(result.get()) != PyTuple_GET_SIZE(grad_input.get())) {
-    char *hook_name = _try_get_name(key, tmp);
+    char *hook_name = _try_get_name(hook, tmp);
     THPUtils_setError("backward hook %s%s%sreturned an incorrect number "
         "of gradients (got %ld, but expected %ld)",
         OPTIONAL_HOOK_NAME,
@@ -534,7 +535,7 @@ static void _ensure_correct_hook_result(THPObjectPtr& grad_input,
   for (int i = 0; i < size; i++) {
     PyObject *original = PyTuple_GET_ITEM(grad_input.get(), i);
     PyObject *returned = PyTuple_GET_ITEM(result.get(), i);
-    _ensure_correct_hook_result_single(original, returned, key);
+    _ensure_correct_hook_result_single(original, returned, hook);
   }
 }
 
@@ -570,7 +571,7 @@ static void _call_output_hooks(THPFunction *self, THPObjectPtr& grad_output)
       if (result.get() != Py_None) {
         // Check all possible inconsistencies of the output that we can detect
         // (sizes, types, etc.)
-        _ensure_correct_hook_result_single(old_grad, result, key);
+        _ensure_correct_hook_result_single(old_grad, result, value);
 
         // Replace the old gradient
         PyTuple_SET_ITEM(new_grad_output.get(), i, result.release());
@@ -603,7 +604,7 @@ static void _call_function_hooks(THPFunction *self, THPObjectPtr& grad_input, TH
       _ensure_tuple(result);
       // Check all possible inconsistencies of the output that we can detect
       // (sizes, types, etc.)
-      _ensure_correct_hook_result(grad_input, result, key);
+      _ensure_correct_hook_result(grad_input, result, value);
       grad_input = result.release();
     }
   }
diff --git a/torch/nn/modules/module.py b/torch/nn/modules/module.py
diff --git a/torch/utils/hooks.py b/torch/utils/hooks.py
