[fx2trt] Add more profiling logging (#78)

Summary:
Pull Request resolved: pytorch/fx2trt#78

This is to analyze more on the lowering latency.

Reviewed By: yyetim, khabinov

Differential Revision: D36352458

fbshipit-source-id: d3f40d02b89b815fd7a0c2ad90ab2ab0a1b24207

Author: singlaiiit

fx/fx2trt.py

@@ -1,4 +1,6 @@
+import logging
 import warnings
+from datetime import datetime
 from typing import Any, Callable, Dict, List, NamedTuple, Optional, Sequence
 
 import numpy
@@ -15,6 +17,10 @@
 from .input_tensor_spec import InputTensorSpec
 from .utils import get_dynamic_dims, LowerPrecision, torch_dtype_to_trt
 
+
+_LOGGER: logging.Logger = logging.getLogger(__name__)
+
+
 TRT_INTERPRETER_CALL_PRE_OBSERVER: Observer[
     Callable[[torch.fx.GraphModule], None]
 ] = Observer("TRT_INTERPRETER_CALL_PRE_OBSERVER")
@@ -179,7 +185,12 @@ def run(
             warnings.warn("Current platform doesn't support fast native fp16!")
 
         self.input_specs_iter = 0
+        run_module_start_time = datetime.now()
         super().run()
+        _LOGGER.info(
+            f"Run Module elapsed time: {datetime.now() - run_module_start_time}"
+        )
+        build_engine_start_time = datetime.now()
 
         self.builder.max_batch_size = max_batch_size
         builder_config = self.builder.create_builder_config()
@@ -227,6 +238,9 @@ def run(
             if builder_config.get_timing_cache()
             else bytearray()
         )
+        _LOGGER.info(
+            f"Build TRT engine elapsed time: {datetime.now() - build_engine_start_time}"
+        )
 
         return TRTInterpreterResult(
             engine, self._input_names, self._output_names, serialized_cache