clean up initization

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

Reviewed By: itomatik

Differential Revision: D39628884

fbshipit-source-id: bb1d5d77eeb965dff675c17a8fbc36e4da4e25cd

d2go/__init__.py

@@ -2,6 +2,8 @@
 # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
 
 
-# Set up custom environment before nearly anything else is imported
-# NOTE: this should be the first import (no not reorder)
-from d2go import initializer, optimizer  # NOQA
+from d2go.initializer import initialize_all
+
+# NOTE: by default a list of initializations will run whenever D2Go is first imported,
+# so that users don't need to do any manual iniitialization other than importing `d2go`.
+initialize_all()

d2go/initializer.py

@@ -2,43 +2,48 @@
 # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
 from mobile_cv.common.misc.oss_utils import fb_overwritable
 
+_INITIALIZED = False
+
+
+def initialize_all():
+    global _INITIALIZED
+    if _INITIALIZED:
+        return
+    _INITIALIZED = True
+
+    _initialize_all()
+
+
+def _initialize_all():
+    _setup_env()
+    _register_builtin_datasets()
+    _populate_registries()
+
+
+# fmt: off
+
 
 @fb_overwritable()
 def _setup_env():
-    # Set up custom environment before nearly anything else is imported
-    # NOTE: this should be the first import (no not reorder)
-    from detectron2.utils.env import (  # noqa F401 isort:skip
-        setup_environment as d2_setup_environment,  # noqa
-    )
+    # register torch vision ops
+    from torchvision.ops import nms  # noqa
+
+    # setup Detectron2 environments
+    from detectron2.utils.env import setup_environment as setup_d2_environment # isort:skip
+    setup_d2_environment()
 
 
 @fb_overwritable()
-def _register_d2_datasets():
-    # this will register D2 builtin datasets
+def _register_builtin_datasets():
+    # Register D2 builtin datasets
     import detectron2.data  # noqa F401
 
 
 @fb_overwritable()
-def _register():
+def _populate_registries():
+    from d2go import optimizer  # noqa
     from d2go.data import dataset_mappers  # noqa
-    from d2go.data.datasets import (  # noqa
-        register_builtin_datasets,  # noqa
-        register_json_datasets,  # noqa
-    )
     from d2go.modeling.backbone import fbnet_v2  # noqa
 
 
-@fb_overwritable()
-def initialize_all():
-    # exclude torch from timing
-    from torchvision.ops import nms  # noqa
-
-    _setup_env()
-    _register_d2_datasets()
-    _register()
-
-
-_INITIALIZED = False
-if not _INITIALIZED:
-    initialize_all()
-    _INITIALIZED = True
+# fmt: on