Implement config validation to find unused keys

Summary:
Implement a `ClassyMap` type which supports tracking reads and freezing the map (the latter is unused currently).

Added it to `ClassificationTask` to catch cases where we don't use any keys passed by users.

This will not catch all instances, like when some components do a deepcopy - we assume all the keys and sub-keys are read in such a situation

Differential Revision: D25321360

fbshipit-source-id: aff06f1b3334ca9d217453d590ac413b4f586966

Author: mannatsingh

classy_vision/configuration/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .classy_map import ClassyMap
+from .config_error import ConfigError, ConfigUnusedKeysError
+
+__all__ = ["ClassyMap", "ConfigError", "ConfigUnusedKeysError"]

classy_vision/configuration/classy_map.py

@@ -0,0 +1,170 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import copy
+from collections.abc import MutableMapping, Mapping
+
+
+class ClassyMap(MutableMapping):
+    """Mapping which can be made immutable. Also supports tracking unused keys."""
+
+    def __init__(self, *args, **kwargs):
+        """Create a ClassyMap.
+
+        Supports the same API as a dict and recursively converts all dicts to
+        ClassyMaps.
+        """
+
+        # NOTE: Another way to implement this would be to subclass dict, but since dict
+        # is a built-in, it isn't treated like a regular MutableMapping, and calls like
+        # func(**map) are handled mysteriously, probably interpreter dependent.
+        # The downside with this implementation is that this isn't a full dict and is
+        # just a mapping, which means some features like JSON serialization don't work
+
+        self._dict = dict(*args, **kwargs)
+        self._frozen = False
+        self._keys_read = set()
+        for k, v in self._dict.items():
+            self._dict[k] = self.to_classy_dict(v)
+
+    def to_classy_dict(self, obj):
+        """Recursively convert all sub items inside obj to ClassyMaps"""
+
+        if isinstance(obj, Mapping):
+            obj = ClassyMap({k: self.to_classy_dict(v) for k, v in obj.items()})
+        elif isinstance(obj, (list, tuple)):
+            # tuples are also converted to lists
+            obj = [self.to_classy_dict(v) for v in obj]
+        return obj
+
+    def keys(self):
+        return self._dict.keys()
+
+    def items(self):
+        self._keys_read.update(self._dict.keys())
+        return self._dict.items()
+
+    def values(self):
+        self._keys_read.update(self._dict.keys())
+        return self._dict.values()
+
+    def pop(self, key, default=None):
+        return self._dict.pop(key, default)
+
+    def popitem(self):
+        return self._dict.popitem()
+
+    def clear(self):
+        self._dict.clear()
+
+    def update(self, *args, **kwargs):
+        if self._frozen:
+            raise TypeError("Frozen ClassyMaps do not support updates")
+        self._dict.update(*args, **kwargs)
+
+    def setdefault(self, key, default=None):
+        return self._dict.setdefault(key, default)
+
+    def __contains__(self, key):
+        return key in self._dict
+
+    def __eq__(self, obj):
+        return self._dict == obj
+
+    def __len__(self):
+        return len(self._dict)
+
+    def __getitem__(self, key):
+        self._keys_read.add(key)
+        return self._dict.__getitem__(key)
+
+    def __iter__(self):
+        return iter(self._dict)
+
+    def __str__(self):
+        return str(self._dict)
+
+    def __repr__(self):
+        return repr(self._dict)
+
+    def get(self, key, default=None):
+        if key in self._dict.keys():
+            self._keys_read.add(key)
+        return self._dict.get(key, default)
+
+    def __copy__(self):
+        ret = ClassyMap()
+        for key, value in self._dict.items():
+            self._keys_read.add(key)
+            ret._dict[key] = value
+
+    def copy(self):
+        return self.__copy__()
+
+    def __deepcopy__(self, memo=None):
+        # for deepcopies we mark all the keys and sub-keys as read
+        ret = ClassyMap()
+        for key, value in self._dict.items():
+            self._keys_read.add(key)
+            ret._dict[key] = copy.deepcopy(value)
+        return ret
+
+    def __setitem__(self, key, value):
+        if self._frozen:
+            raise TypeError("Frozen ClassyMaps do not support assignment")
+        if isinstance(value, dict) and not isinstance(value, ClassyMap):
+            value = ClassyMap(value)
+        self._dict.__setitem__(key, value)
+
+    def __delitem__(self, key):
+        if self._frozen:
+            raise TypeError("Frozen ClassyMaps do not support key deletion")
+        del self._dict[key]
+
+    def _freeze(self, obj):
+        if isinstance(obj, Mapping):
+            assert isinstance(obj, ClassyMap), f"{obj} is not a ClassyMap"
+            obj._frozen = True
+            for value in obj.values():
+                self._freeze(value)
+        elif isinstance(obj, list):
+            for value in obj:
+                self._freeze(value)
+
+    def _reset_tracking(self, obj):
+        if isinstance(obj, Mapping):
+            assert isinstance(obj, ClassyMap), f"{obj} is not a ClassyMap"
+            obj._keys_read = set()
+            for value in obj._dict.values():
+                self._reset_tracking(value)
+        elif isinstance(obj, list):
+            for value in obj:
+                self._reset_tracking(value)
+
+    def _unused_keys(self, obj):
+        unused_keys = []
+        if isinstance(obj, Mapping):
+            assert isinstance(obj, ClassyMap), f"{obj} is not a ClassyMap"
+            unused_keys = [key for key in obj._dict.keys() if key not in obj._keys_read]
+            for key, value in obj._dict.items():
+                unused_keys += [
+                    f"{key}.{subkey}" for subkey in self._unused_keys(value)
+                ]
+        elif isinstance(obj, list):
+            for i, value in enumerate(obj):
+                unused_keys += [f"{i}.{subkey}" for subkey in self._unused_keys(value)]
+        return unused_keys
+
+    def freeze(self):
+        """Freeze the ClassyMap to disallow mutations"""
+        self._freeze(self)
+
+    def reset_tracking(self):
+        """Reset key tracking"""
+        self._reset_tracking(self)
+
+    def unused_keys(self):
+        """Fetch all the unused keys"""
+        return self._unused_keys(self)

classy_vision/configuration/config_error.py

@@ -0,0 +1,16 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import List
+
+
+class ConfigError(Exception):
+    pass
+
+
+class ConfigUnusedKeysError(ConfigError):
+    def __init__(self, unused_keys: List[str]):
+        self.unused_keys = unused_keys
+        super().__init__(f"The following keys were unused: {self.unused_keys}")

classy_vision/optim/sgd.py

@@ -7,6 +7,7 @@
 from typing import Any, Dict
 
 import torch.optim
+from classy_vision.configuration import ClassyMap
 
 from . import ClassyOptimizer, register_optimizer
 
@@ -63,10 +64,11 @@ def from_config(cls, config: Dict[str, Any]) -> "SGD":
         config.setdefault("weight_decay", 0.0)
         config.setdefault("nesterov", False)
         config.setdefault("use_larc", False)
-        config.setdefault(
-            "larc_config", {"clip": True, "eps": 1e-08, "trust_coefficient": 0.02}
-        )
-
+        if config["use_larc"]:
+            larc_config = ClassyMap(clip=True, eps=1e-8, trust_coefficient=0.02)
+        else:
+            larc_config = None
+        config.setdefault("larc_config", larc_config)
         assert (
             config["momentum"] >= 0.0
             and config["momentum"] < 1.0

classy_vision/tasks/classification_task.py

@@ -16,6 +16,7 @@
 
 import torch
 import torch.nn as nn
+from classy_vision.configuration import ClassyMap, ConfigUnusedKeysError
 from classy_vision.dataset import ClassyDataset, build_dataset
 from classy_vision.dataset.transforms.mixup import MixupTransform
 from classy_vision.generic.distributed_util import (
@@ -456,6 +457,11 @@ def from_config(cls, config: Dict[str, Any]) -> "ClassificationTask":
         Returns:
             A ClassificationTask instance.
         """
+        orig_config = config
+        config = ClassyMap(orig_config)
+        # access the name key to make sure it gets tracked
+        config["name"]
+
         test_only = config.get("test_only", False)
         if not test_only:
             # TODO Make distinction between epochs and phases in optimizer clear
@@ -537,9 +543,13 @@ def from_config(cls, config: Dict[str, Any]) -> "ClassificationTask":
         for phase_type in datasets:
             task.set_dataset(datasets[phase_type], phase_type)
 
+        # at this stage all the configs keys should have been used
+        if config.unused_keys():
+            raise ConfigUnusedKeysError(config.unused_keys())
+
         # NOTE: this is a private member and only meant to be used for
         # logging/debugging purposes. See __repr__ implementation
-        task._config = config
+        task._config = orig_config
 
         return task
 

test/configuration_classy_map_test.py

@@ -0,0 +1,101 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import copy
+import unittest
+
+from classy_vision.configuration import ClassyMap
+
+
+class ClassyMapTest(unittest.TestCase):
+    def test_dict(self):
+        d = ClassyMap(a=1, b=[1, 2, "3"])
+        d["c"] = [4]
+        d["d"] = {"a": 2}
+        self.assertEqual(d, {"a": 1, "b": [1, 2, "3"], "c": [4], "d": {"a": 2}})
+        self.assertIsInstance(d, ClassyMap)
+        self.assertIsInstance(d["d"], ClassyMap)
+
+    def test_freezing(self):
+        d = ClassyMap(a=1, b=2)
+        d.freeze()
+        # resetting an already existing key
+        with self.assertRaises(TypeError):
+            d["a"] = 3
+        # adding a new key
+        with self.assertRaises(TypeError):
+            d["f"] = 3
+
+    def test_unused_keys(self):
+        d = ClassyMap(
+            a=1,
+            b=[
+                1,
+                2,
+                {
+                    "c": {"a": 2},
+                    "d": 4,
+                    "e": {"a": 1, "b": 2},
+                    "f": {"a": 1, "b": {"c": 2}},
+                },
+            ],
+        )
+
+        all_keys = {
+            "a",
+            "b",
+            "b.2.c",
+            "b.2.c.a",
+            "b.2.d",
+            "b.2.e",
+            "b.2.f",
+            "b.2.e.a",
+            "b.2.e.b",
+            "b.2.f.a",
+            "b.2.f.b",
+            "b.2.f.b.c",
+        }
+
+        def test_func(**kwargs):
+            return None
+
+        for _ in range(2):
+            expected_unused_keys = all_keys.copy()
+            self.assertSetEqual(set(d.unused_keys()), expected_unused_keys)
+
+            _ = d["a"]
+            expected_unused_keys.remove("a")
+            self.assertSetEqual(set(d.unused_keys()), expected_unused_keys)
+
+            _ = d["b"][2].get("d")
+            expected_unused_keys.remove("b")
+            expected_unused_keys.remove("b.2.d")
+            self.assertSetEqual(set(d.unused_keys()), expected_unused_keys)
+
+            _ = d["b"][2]["e"]
+            expected_unused_keys.remove("b.2.e")
+            self.assertSetEqual(set(d.unused_keys()), expected_unused_keys)
+
+            _ = d["b"][2]["e"].items()
+            expected_unused_keys.remove("b.2.e.a")
+            expected_unused_keys.remove("b.2.e.b")
+            self.assertSetEqual(set(d.unused_keys()), expected_unused_keys)
+
+            _ = d["b"][2]["f"]
+            expected_unused_keys.remove("b.2.f")
+            self.assertSetEqual(set(d.unused_keys()), expected_unused_keys)
+
+            test_func(**d["b"][2]["f"])
+            expected_unused_keys.remove("b.2.f.a")
+            expected_unused_keys.remove("b.2.f.b")
+            self.assertSetEqual(set(d.unused_keys()), expected_unused_keys)
+
+            _ = copy.deepcopy(d)
+            expected_unused_keys.remove("b.2.c")
+            expected_unused_keys.remove("b.2.c.a")
+            expected_unused_keys.remove("b.2.f.b.c")
+            self.assertSetEqual(set(d.unused_keys()), expected_unused_keys)
+
+            d.reset_tracking()