implement nn.Module.__dir__ (pytorch#1142)

Author: chenyuntc

test/test_nn.py

@@ -467,6 +467,19 @@ def test_children(self):
         s = nn.Sequential(l1, l2, l1, l2, subnet)
         self.assertEqual(list(s.children()), [l1, l2, subnet])
 
+    def test_dir(self):
+        linear = nn.Linear(2, 2)
+        linear._test_submodule = nn.Linear(2, 2)
+        linear._test_parameter = Parameter(torch.Tensor(2, 2))
+        linear.register_buffer('_test_buffer', torch.Tensor(2, 2))
+        keys = linear.__dir__()
+        self.assertIn('_test_submodule', keys)
+        self.assertIn('_test_parameter', keys)
+        self.assertIn('_test_buffer', keys)
+
+        for key in keys:
+            self.assertTrue(hasattr(linear, key))
+
     def test_named_children(self):
         l1 = nn.Linear(2, 2)
         l2 = nn.Linear(2, 2)

torch/nn/modules/module.py

@@ -462,3 +462,12 @@ def __repr__(self):
             tmpstr = tmpstr + '  (' + key + '): ' + modstr + '\n'
         tmpstr = tmpstr + ')'
         return tmpstr
+
+    def __dir__(self):
+        module_attrs = dir(self.__class__)
+        attrs = list(self.__dict__.keys())
+        parameters = list(self._parameters.keys())
+        modules = list(self._modules.keys())
+        buffers = list(self._buffers.keys())
+        keys = module_attrs + attrs + parameters + modules + buffers
+        return sorted(keys)