Add Module.requires_grad_ (#22576)

ssnl · facebook-github-bot · commit 478d480d376f · 2019-07-08T12:13:07.000-07:00
Summary: addresses #20241 Pull Request resolved: #22576 Differential Revision: D16149314 Pulled By: zou3519 fbshipit-source-id: 1cc4c1ec084df30e00e9ae73ce1a53494a034d5c
diff --git a/test/test_nn.py b/test/test_nn.py
@@ -570,6 +570,19 @@ def __init__(self):
 
         return l, n, s
 
+    def test_requires_grad_(self):
+        m = self._create_basic_net()[-1]
+        assert len(list(m.buffers())) > 0, 'invalid test'
+        assert all(not b.requires_grad for b in m.buffers()) > 0, 'invalid test'
+        assert len(list(m.parameters())) > 0, 'invalid test'
+        assert all(p.requires_grad for p in m.parameters()) > 0, 'invalid test'
+        for requires_grad in (False, True):
+            self.assertIs(m.requires_grad_(requires_grad), m)
+            for p in m.parameters():
+                self.assertEqual(p.requires_grad, requires_grad)
+            for b in m.buffers():
+                self.assertFalse(b.requires_grad)
+
     def test_module_backcompat(self):
         from torch.serialization import SourceChangeWarning
         path = download_file('https://download.pytorch.org/test_data/linear.pt')
diff --git a/torch/nn/modules/module.py b/torch/nn/modules/module.py
@@ -1050,6 +1050,10 @@ def train(self, mode=True):
         mode, if they are affected, e.g. :class:`Dropout`, :class:`BatchNorm`,
         etc.
 
+        Args:
+            mode (bool): whether to set training mode (``True``) or evaluation
+                         mode (``False``). Default: ``True``.
+
         Returns:
             Module: self
         """
@@ -1065,9 +1069,35 @@ def eval(self):
         particular modules for details of their behaviors in training/evaluation
         mode, if they are affected, e.g. :class:`Dropout`, :class:`BatchNorm`,
         etc.
+
+        This is equivalent with :meth:`self.train(False) <torch.nn.Module.train>`.
+
+        Returns:
+            Module: self
         """
         return self.train(False)
 
+    def requires_grad_(self, requires_grad=True):
+        r"""Change if autograd should record operations on parameters in this
+        module.
+
+        This method sets the parameters' :attr:`requires_grad` attributes
+        in-place.
+
+        This method is helpful for freezing part of the module for finetuning
+        or training parts of a model individually (e.g., GAN training).
+
+        Args:
+            requires_grad (bool): whether autograd should record operations on
+                                  parameters in this module. Default: ``True``.
+
+        Returns:
+            Module: self
+        """
+        for p in self.parameters():
+            p.requires_grad_(requires_grad)
+        return self
+
     def zero_grad(self):
         r"""Sets gradients of all model parameters to zero."""
         for p in self.parameters():
