Integrating model profiler into d2go NAS task

Reviewed By: larryliu0820

Differential Revision: D30390706

fbshipit-source-id: 49f83f884f497df227448f7e59903bd1bd6e5484

d2go/export/api.py

@@ -88,21 +88,12 @@ class PredictorExportConfig(NamedTuple):
     run_func_info: FuncInfo = FuncInfo.gen_func_info(NaiveRunFunc, params={})
 
 
-def convert_and_export_predictor(
+def convert_predictor(
     cfg,
     pytorch_model,
     predictor_type,
-    output_dir,
     data_loader,
 ):
-    """
-    Entry point for convert and export model. This involves two steps:
-        - convert: converting the given `pytorch_model` to another format, currently
-            mainly for quantizing the model.
-        - export: exporting the converted `pytorch_model` to predictor. This step
-            should not alter the behaviour of model.
-    """
-
     if "int8" in predictor_type:
         if not cfg.QUANTIZATION.QAT.ENABLED:
             logger.info(
@@ -131,7 +122,24 @@ def convert_and_export_predictor(
         logger.info("Fused Model:\n{}".format(pytorch_model))
         if fuse_utils.count_bn_exist(pytorch_model) > 0:
             logger.warning("BN existed in pytorch model after fusing.")
+    return pytorch_model
 
+
+def convert_and_export_predictor(
+    cfg,
+    pytorch_model,
+    predictor_type,
+    output_dir,
+    data_loader,
+):
+    """
+    Entry point for convert and export model. This involves two steps:
+        - convert: converting the given `pytorch_model` to another format, currently
+            mainly for quantizing the model.
+        - export: exporting the converted `pytorch_model` to predictor. This step
+            should not alter the behaviour of model.
+    """
+    pytorch_model = convert_predictor(cfg, pytorch_model, predictor_type, data_loader)
     return export_predictor(cfg, pytorch_model, predictor_type, output_dir, data_loader)
 
 

d2go/export/torchscript.py

@@ -140,15 +140,7 @@ def tracing_adapter_wrap_export(old_f):
     return new_f
 
 
-def tracing_adapter_wrap_load(old_f):
-    def new_f(cls, save_path, **load_kwargs):
-        assert "inputs_schema" in load_kwargs, load_kwargs.keys()
-        assert "outputs_schema" in load_kwargs, load_kwargs.keys()
-        inputs_schema = instantiate(load_kwargs.pop("inputs_schema"))
-        outputs_schema = instantiate(load_kwargs.pop("outputs_schema"))
-        traced_model = old_f(cls, save_path, **load_kwargs)
-
-        class TracingAdapterModelWrapper(nn.Module):
+class TracingAdapterModelWrapper(nn.Module):
     def __init__(self, traced_model, inputs_schema, outputs_schema):
         super().__init__()
         self.traced_model = traced_model
@@ -160,6 +152,15 @@ def tracing_adapter_wrap_load(old_f):
         flattened_outputs = self.traced_model(*flattened_inputs)
         return self.outputs_schema(flattened_outputs)
 
+
+def tracing_adapter_wrap_load(old_f):
+    def new_f(cls, save_path, **load_kwargs):
+        assert "inputs_schema" in load_kwargs, load_kwargs.keys()
+        assert "outputs_schema" in load_kwargs, load_kwargs.keys()
+        inputs_schema = instantiate(load_kwargs.pop("inputs_schema"))
+        outputs_schema = instantiate(load_kwargs.pop("outputs_schema"))
+        traced_model = old_f(cls, save_path, **load_kwargs)
+
         return TracingAdapterModelWrapper(traced_model, inputs_schema, outputs_schema)
 
     return new_f

d2go/modeling/meta_arch/rcnn.py

@@ -62,11 +62,16 @@ class GeneralizedRCNNPatch:
 def default_rcnn_prepare_for_export(self, cfg, inputs, predictor_type):
 
     if "torchscript" in predictor_type and "@tracing" in predictor_type:
+        preprocess_info = FuncInfo.gen_func_info(
+            D2RCNNTracingWrapper.Preprocess, params={}
+        )
+        preprocess_func = preprocess_info.instantiate()
         return PredictorExportConfig(
             model=D2RCNNTracingWrapper(self),
-            data_generator=D2RCNNTracingWrapper.generator_trace_inputs,
-            run_func_info=FuncInfo.gen_func_info(
-                D2RCNNTracingWrapper.RunFunc, params={}
+            data_generator=lambda x: (preprocess_func(x),),
+            preprocess_info=preprocess_info,
+            postprocess_info=FuncInfo.gen_func_info(
+                D2RCNNTracingWrapper.Postprocess, params={}
             ),
         )
 

d2go/utils/export_utils.py

@@ -100,15 +100,18 @@ class D2RCNNTracingWrapper(nn.Module):
         return self.model.inference(inputs, do_postprocess=False)[0]
 
     @staticmethod
-    def generator_trace_inputs(batch):
+    class Preprocess(object):
         """
         This function describes how to covert orginal input (from the data loader)
         to the inputs used during the tracing (i.e. the inputs of forward function).
         """
-        return (batch[0]["image"],)
 
-    class RunFunc(object):
-        def __call__(self, tracing_adapter_wrapper, batch):
+        def __call__(self, batch):
+            assert len(batch) == 1, "only support single batch"
+            return batch[0]["image"]
+
+    class Postprocess(object):
+        def __call__(self, batch, inputs, outputs):
             """
             This function describes how to run the predictor using exported model. Note
             that `tracing_adapter_wrapper` runs the traced model under the hood and
@@ -116,7 +119,5 @@ class D2RCNNTracingWrapper(nn.Module):
             """
             assert len(batch) == 1, "only support single batch"
             width, height = batch[0]["width"], batch[0]["height"]
-            inputs = D2RCNNTracingWrapper.generator_trace_inputs(batch)
-            results_per_image = tracing_adapter_wrapper(inputs)
-            r = detector_postprocess(results_per_image, height, width)
+            r = detector_postprocess(outputs, height, width)
             return [{"instances": r}]