refactor(//py): Have the python api default input types to PyTorch like

behavior unless user explicitly overrides

Signed-off-by: Naren Dasan <[email protected]>
Signed-off-by: Naren Dasan <[email protected]>

Author: narendasan

py/setup.py

@@ -181,7 +181,7 @@ def run(self):
         include_dirs=[
             dir_path + "trtorch/csrc",
             dir_path + "/../",
-            dir_path + "/../bazel-trtorch-testing/external/tensorrt/include",
+            dir_path + "/../bazel-TRTorch/external/tensorrt/include",
         ],
         extra_compile_args=[
             "-Wno-deprecated",

py/trtorch/Input.py

@@ -30,6 +30,7 @@ class _ShapeMode(Enum):
     shape_mode = None
     shape = None
     dtype = _types.dtype.float32
+    _explicit_set_dtype = False
     format = _types.TensorFormat.contiguous
 
     def __init__(self, *args, **kwargs):
@@ -105,6 +106,7 @@ def __init__(self, *args, **kwargs):
 
         if "dtype" in kwargs:
             self.dtype = Input._parse_dtype(kwargs["dtype"])
+            self._explicit_set_dtype = True
 
         if "format" in kwargs:
             self.format = Input._parse_format(kwargs["format"])
@@ -128,6 +130,7 @@ def _to_internal(self) -> trtorch._C.Input:
             internal_in.opt = self.shape
             internal_in.input_is_dynamic = False
         internal_in.dtype = self.dtype
+        internal_in._explicit_set_dtype = self._explicit_set_dtype
         internal_in.format = self.format
         return internal_in
 

py/trtorch/_compile_spec.py

@@ -169,6 +169,16 @@ def _parse_compile_spec(compile_spec: Dict[str, Any]) -> trtorch._C.CompileSpec:
 
     if "enabled_precisions" in compile_spec:
         info.enabled_precisions = _parse_enabled_precisions(compile_spec["enabled_precisions"])
+        # We want default behavior to match PyTorch, so in the case the user did not explicitly set the dtype for inputs they
+        # will follow PyTorch convetions
+        for i in info.inputs:
+            if not i._explicit_set_dtype:
+                if _types.dtype.int8 in info.enabled_precisions:
+                    i.dtype = _types.dtype.float32
+                elif _types.dtype.half in info.enabled_precisions:
+                    i.dtype = _types.dtype.float16
+                else:
+                    i.dtype = _types.dtype.float32
 
     if "calibrator" in compile_spec:
         info.ptq_calibrator = compile_spec["calibrator"]

py/trtorch/csrc/tensorrt_classes.h

@@ -41,13 +41,15 @@ struct Input : torch::CustomClassHolder {
   std::vector<int64_t> max;
 
   bool input_is_dynamic;
+  bool explicit_set_dtype;
   DataType dtype;
   TensorFormat format;
 
   ADD_FIELD_GET_SET(min, std::vector<int64_t>);
   ADD_FIELD_GET_SET(opt, std::vector<int64_t>);
   ADD_FIELD_GET_SET(max, std::vector<int64_t>);
   ADD_FIELD_GET_SET(input_is_dynamic, bool);
+  ADD_FIELD_GET_SET(explicit_set_dtype, bool);
   ADD_ENUM_GET_SET(dtype, DataType, static_cast<int64_t>(DataType::kBool));
   ADD_ENUM_GET_SET(format, TensorFormat, static_cast<int64_t>(TensorFormat::kContiguous));
 

py/trtorch/csrc/trtorch_py.cpp

@@ -170,6 +170,7 @@ PYBIND11_MODULE(_C, m) {
       .def_readwrite("opt", &Input::opt)
       .def_readwrite("max", &Input::max)
       .def_readwrite("input_is_dynamic", &Input::input_is_dynamic)
+      .def_readwrite("_explicit_set_dtype", &Input::explicit_set_dtype)
       .def_readwrite("dtype", &Input::dtype)
       .def_readwrite("format", &Input::format);