Integrate the Genie optimization engine to d2go (reapply D47502855)

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

Genie optimization engine has the assumption that when a training iteration is started, it is also finished. And the after_step hook is called. This assumption is not valid in d2go.
https://www.internalfb.com/code/fbsource/[1537eddbd235e3f599709a493c1a80c7d016b3f8]/fbcode/vision/fair/detectron2/detectron2/engine/train_loop.py?lines=151-165

When an exception is triggered, the last iteration's after_step hook is not called.

In this diff, we patch up the hook integration to ensure that the Genie after_step hook is always called.

everything else remain the same as  D47502855

Reviewed By: XiaoliangDai

Differential Revision: D47611143

fbshipit-source-id: b8b1ae2f304a40cf74340bbaf35647332a9a1524

d2go/config/__init__.py

@@ -8,6 +8,7 @@ from d2go.config.config import (
     CfgNode,
     CONFIG_CUSTOM_PARSE_REGISTRY,
     CONFIG_SCALING_METHOD_REGISTRY,
+    convert_cfg_to_dict,
     load_full_config_from_file,
     temp_defrost,
     temp_new_allowed,
@@ -20,6 +21,7 @@ __all__ = [
     "CONFIG_SCALING_METHOD_REGISTRY",
     "CfgNode",
     "auto_scale_world_size",
+    "convert_cfg_to_dict",
     "load_full_config_from_file",
     "reroute_config_path",
     "temp_defrost",

d2go/config/config.py

@@ -190,3 +190,13 @@ def load_full_config_from_file(filename: str) -> CfgNode:
     cfg = loaded_cfg.get_default_cfg()
     cfg.merge_from_other_cfg(loaded_cfg)
     return cfg
+
+
+def convert_cfg_to_dict(cfg):
+    if not isinstance(cfg, CfgNode):
+        return cfg
+    else:
+        cfg_dict = dict(cfg)
+        for k, v in cfg_dict.items():
+            cfg_dict[k] = convert_cfg_to_dict(v)
+        return cfg_dict