Parse SuperNet search space from config

Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/168

Add a hook in D2 (https://github.com/facebookresearch/d2go/commit/7992f91324aee6ae59795063a007c6837e60cdb8)Go config for custom parsing so we can support custom objects in D2 (https://github.com/facebookresearch/d2go/commit/7992f91324aee6ae59795063a007c6837e60cdb8)Go config like the search space objects.
Then adding SuperNet custom config processing to parse search space from arch_def when supernet is enabled, so it can be used in D2 (https://github.com/facebookresearch/d2go/commit/7992f91324aee6ae59795063a007c6837e60cdb8)Go SuperNet training.

This is an alternative approach to D33191150. In this approach we parse the entire architecture as a search space which will not have the limitations that we have in parsing only the dynamic blocks parts.

Reviewed By: zhanghang1989

Differential Revision: D33793423

fbshipit-source-id: 8acf5c5afb3c5c0005bdb0ca16847026e1b45e2c

d2go/config/__init__.py

@@ -4,6 +4,7 @@
 
 # forward the namespace to avoid `d2go.config.config`
 from .config import (
+    CONFIG_CUSTOM_PARSE_REGISTRY,
     CONFIG_SCALING_METHOD_REGISTRY,
     CfgNode,
     auto_scale_world_size,
@@ -13,6 +14,7 @@ from .config import (
 
 
 __all__ = [
+    "CONFIG_CUSTOM_PARSE_REGISTRY",
     "CONFIG_SCALING_METHOD_REGISTRY",
     "CfgNode",
     "auto_scale_world_size",

d2go/config/config.py

@@ -4,6 +4,7 @@
 
 import contextlib
 import logging
+from typing import List
 
 import mock
 import yaml
@@ -14,6 +15,8 @@ from .utils import reroute_config_path
 
 logger = logging.getLogger(__name__)
 
+CONFIG_CUSTOM_PARSE_REGISTRY = Registry("CONFIG_CUSTOM_PARSE")
+
 
 class CfgNode(_CfgNode):
     @classmethod
@@ -28,7 +31,18 @@ class CfgNode(_CfgNode):
     def merge_from_file(self, cfg_filename: str, *args, **kwargs):
         cfg_filename = reroute_config_path(cfg_filename)
         with reroute_load_yaml_with_base():
-            return super().merge_from_file(cfg_filename, *args, **kwargs)
+            res = super().merge_from_file(cfg_filename, *args, **kwargs)
+            self._run_custom_processing(is_dump=False)
+            return res
+
+    def merge_from_list(self, cfg_list: List[str]):
+        res = super().merge_from_list(cfg_list)
+        self._run_custom_processing(is_dump=False)
+        return res
+
+    def dump(self, *args, **kwargs):
+        self._run_custom_processing(is_dump=True)
+        return super().dump(*args, **kwargs)
 
     @staticmethod
     def load_yaml_with_base(filename: str, *args, **kwargs):
@@ -39,6 +53,16 @@ class CfgNode(_CfgNode):
         # dump follows alphabetical order, thus good for hash use
         return hash(self.dump())
 
+    def _run_custom_processing(self, is_dump=False):
+        """Apply config load post custom processing from registry"""
+        frozen = self.is_frozen()
+        self.defrost()
+        for name, process_func in CONFIG_CUSTOM_PARSE_REGISTRY:
+            logger.info(f"Apply config processing: {name}, is_dump={is_dump}")
+            process_func(self, is_dump)
+        if frozen:
+            self.freeze()
+
 
 @contextlib.contextmanager
 def temp_defrost(cfg):