Add model parameters / modules inspection helper. (#2466)

* Adds inspection method on models to view parameters and modules.

* Add a test.

* Fix a typo.

* Fix another typo.

* Fix mypy and docs.

* Revert changes on model.

* Add inspect_model_parameters and tests.

* Avert circular import (for now).

* Remove extra blank lines.

* Allow modules instead of restricting to only model.

* Add a blank line.

* Allow too many lines in util_test.py

* update inspection util.

* pylint and mypy.

* Update docstring.

allennlp/nn/util.py

@@ -7,6 +7,7 @@
 import logging
 import copy
 import math
+import json
 
 import torch
 
@@ -1378,3 +1379,37 @@ def uncombine_initial_dims(tensor: torch.Tensor, original_size: torch.Size) -> t
     else:
         view_args = list(original_size) + [tensor.size(-1)]
         return tensor.view(*view_args)
+
+
+def inspect_parameters(module: torch.nn.Module, quiet: bool = False) -> Dict[str, Any]:
+    """
+    Inspects the model/module parameters and their tunability. The output is structured
+    in a nested dict so that parameters in same sub-modules are grouped together.
+    This can be helpful to setup module path based regex, for example in initializer.
+    It prints it by default (optional) and returns the inspection dict. Eg. output::
+
+        {
+            "_text_field_embedder": {
+                "token_embedder_tokens": {
+                    "_projection": {
+                        "bias": "tunable",
+                        "weight": "tunable"
+                    },
+                    "weight": "frozen"
+                }
+            }
+        }
+
+    """
+    results: Dict[str, Any] = {}
+    for name, param in sorted(module.named_parameters()):
+        keys = name.split(".")
+        write_to = results
+        for key in keys[:-1]:
+            if key not in write_to:
+                write_to[key] = {}
+            write_to = write_to[key]
+        write_to[keys[-1]] = "tunable" if param.requires_grad else "frozen"
+    if not quiet:
+        print(json.dumps(results, indent=4))
+    return results

allennlp/tests/fixtures/decomposable_attention/parameters_inspection.json

@@ -0,0 +1,39 @@
+{
+    "_aggregate_feedforward": {
+        "_linear_layers": {
+            "0": {
+                "bias": "tunable",
+                "weight": "tunable"
+            }
+        }
+    },
+    "_attend_feedforward": {
+        "_module": {
+            "_linear_layers": {
+                "0": {
+                    "bias": "tunable",
+                    "weight": "tunable"
+                }
+            }
+        }
+    },
+    "_compare_feedforward": {
+        "_module": {
+            "_linear_layers": {
+                "0": {
+                    "bias": "tunable",
+                    "weight": "tunable"
+                }
+            }
+        }
+    },
+    "_text_field_embedder": {
+        "token_embedder_tokens": {
+            "_projection": {
+                "bias": "tunable",
+                "weight": "tunable"
+            },
+            "weight": "frozen"
+        }
+    }
+}

allennlp/tests/nn/util_test.py

@@ -1,4 +1,5 @@
 # pylint: disable=invalid-name,no-self-use,too-many-public-methods,not-callable,too-many-lines,protected-access
+import json
 from typing import NamedTuple
 
 import numpy
@@ -9,6 +10,7 @@
 from allennlp.common.checks import ConfigurationError
 from allennlp.common.testing import AllenNlpTestCase
 from allennlp.nn import util
+from allennlp.models import load_archive
 
 
 class TestNnUtil(AllenNlpTestCase):
@@ -1023,6 +1025,14 @@ def test_uncombine_initial_dims(self):
         embedding = util.uncombine_initial_dims(embedding2d, torch.Size((4, 10, 20, 17, 5)))
         assert list(embedding.size()) == [4, 10, 20, 17, 5, 12]
 
+    def test_inspect_model_parameters(self):
+        model_archive = str(self.FIXTURES_ROOT / 'decomposable_attention' / 'serialization' / 'model.tar.gz')
+        parameters_inspection = str(self.FIXTURES_ROOT / 'decomposable_attention' / 'parameters_inspection.json')
+        model = load_archive(model_archive).model
+        with open(parameters_inspection) as file:
+            parameters_inspection_dict = json.load(file)
+        assert parameters_inspection_dict == util.inspect_parameters(model)
+
     def test_move_to_device(self):
         # We're faking the tensor here so that we can test the calls to .cuda() without actually
         # needing a GPU.