feat(//py): API now produces valid engines that are consumable by

TensorRT Python API

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

Author: narendasan

py/setup.py

@@ -1,5 +1,6 @@
 import os
 import sys
+import glob
 import setuptools
 from setuptools import setup, Extension, find_packages
 from setuptools.command.build_ext import build_ext
@@ -8,7 +9,7 @@
 from distutils.cmd import Command
 
 from torch.utils import cpp_extension
-from shutil import copyfile
+from shutil import copyfile, rmtree
 
 dir_path = os.path.dirname(os.path.realpath(__file__))
 
@@ -60,7 +61,7 @@ def run(self):
 
 class CleanCommand(Command):
     """Custom clean command to tidy up the project root."""
-    PY_CLEAN_FILES = ['./build', './dist', './trtorch/__pycache__', './*.pyc', './*.tgz', './*.egg-info']
+    PY_CLEAN_FILES = ['./build', './dist', './trtorch/__pycache__', './trtorch/lib', './*.pyc', './*.tgz', './*.egg-info']
     description = "Command to tidy up the project root"
     user_options = []
 
@@ -75,11 +76,11 @@ def run(self):
             # Make paths absolute and relative to this path
             abs_paths = glob.glob(os.path.normpath(os.path.join(dir_path, path_spec)))
             for path in [str(p) for p in abs_paths]:
-                if not path.startswith(root_dir):
+                if not path.startswith(dir_path):
                     # Die if path in CLEAN_FILES is absolute + outside this directory
-                    raise ValueError("%s is not a path inside %s" % (path, root_dir))
+                    raise ValueError("%s is not a path inside %s" % (path, dir_path))
                 print('Removing %s' % os.path.relpath(path))
-                shutil.rmtree(path)
+                rmtree(path)
 
 ext_modules = [
     cpp_extension.CUDAExtension('trtorch._C',

py/trtorch/compiler.py

@@ -15,7 +15,7 @@ def _supported_input_size_type(input_size: Any) -> bool:
     else:
         raise TypeError("Input sizes for inputs are required to be a List, tuple or torch.Size or a Dict of three sizes (min, opt, max), found type: " + str(type(input_size)))
 
-def _parse_input_sizes(input_sizes: List) -> List:
+def _parse_input_ranges(input_sizes: List) -> List:
 
     if any (not isinstance(i, dict) and not _supported_input_size_type(i) for i in input_sizes):
         raise KeyError("An input size must either be a static size or a range of three sizes (min, opt, max) as Dict")
@@ -28,16 +28,14 @@ def _parse_input_sizes(input_sizes: List) -> List:
                 in_range.min = i["min"]
                 in_range.opt = i["opt"]
                 in_range.max = i["max"]
-
-                parsed_input_sizes.append(in_range.to_internal_input_range())
+                parsed_input_sizes.append(in_range)
 
             elif "opt" in i:
                 in_range = trtorch._C.InputRange()
                 in_range.min = i["opt"]
                 in_range.opt = i["opt"]
                 in_range.max = i["opt"]
-
-                parsed_input_sizes.append(in_range.to_internal_input_range())
+                parsed_input_sizes.append(in_range)
 
             else:
                 raise KeyError("An input size must either be a static size or a range of three sizes (min, opt, max) as Dict")
@@ -47,8 +45,14 @@ def _parse_input_sizes(input_sizes: List) -> List:
             in_range.min = i
             in_range.opt = i
             in_range.max = i
+            parsed_input_sizes.append(in_range)
 
-            parsed_input_sizes.append(in_range.to_internal_input_range())
+        elif isinstance(i, tuple):
+            in_range = trtorch._C.InputRange()
+            in_range.min = list(i)
+            in_range.opt = list(i)
+            in_range.max = list(i)
+            parsed_input_sizes.append(in_range)
 
     return parsed_input_sizes
 
@@ -87,7 +91,8 @@ def _parse_extra_info(extra_info: Dict[str, Any]) -> trtorch._C._ExtraInfo:
     if "input_shapes" not in extra_info and not isinstance(extra_info["input_shapes"], list):
         raise KeyError("Input shapes for inputs are required as a List, provided as either a static sizes or a range of three sizes (min, opt, max) as Dict")
 
-    info.input_ranges = _parse_input_sizes(extra_info["input_shapes"])
+    info.input_ranges = _parse_input_ranges(extra_info["input_shapes"])
+    print(info.input_ranges)
 
     if "op_precision" in extra_info:
         info.op_precision = _parse_op_precision(extra_info["op_precision"])

py/trtorch/csrc/trtorch_py.cpp

@@ -18,6 +18,9 @@ struct InputRange {
   std::vector<int64_t> max;
 
   core::conversion::InputRange toInternalInputRange() {
+    for (auto o : opt) {
+      std::cout << o << std::endl;
+    }
     return core::conversion::InputRange(min, opt, max);
   }
 };
@@ -76,7 +79,11 @@ nvinfer1::EngineCapability toTRTEngineCapability(EngineCapability value) {
 struct ExtraInfo {
 
   core::ExtraInfo toInternalExtraInfo() {
-    auto info = core::ExtraInfo(input_ranges);
+    std::cout << "HELLO" << input_ranges.size() << std::endl;
+    for (auto i : input_ranges) {
+      internal_input_ranges.push_back(i.toInternalInputRange());
+    }
+    auto info = core::ExtraInfo(internal_input_ranges);
     info.convert_info.engine_settings.op_precision = toTRTDataType(op_precision);
     info.convert_info.engine_settings.refit = refit;
     info.convert_info.engine_settings.debug = debug;
@@ -91,7 +98,8 @@ struct ExtraInfo {
     return info;
   }
 
-  std::vector<core::conversion::InputRange> input_ranges;
+  std::vector<InputRange> input_ranges;
+  std::vector<core::conversion::InputRange> internal_input_ranges;
   DataType op_precision = DataType::kFloat;
   bool refit = false;
   bool debug = false;
@@ -112,10 +120,10 @@ torch::jit::Module CompileGraph(const torch::jit::Module& mod, ExtraInfo& info)
   return trt_mod;
 }
 
-std::string ConvertGraphToTRTEngine(const torch::jit::Module& mod, const std::string& method_name, ExtraInfo& info) {
+py::bytes ConvertGraphToTRTEngine(const torch::jit::Module& mod, const std::string& method_name, ExtraInfo& info) {
   py::gil_scoped_acquire gil;
   auto trt_engine = core::ConvertGraphToTRTEngine(mod, method_name, info.toInternalExtraInfo());
-  return trt_engine;
+  return py::bytes(trt_engine);
 }
 
 bool CheckMethodOperatorSupport(const torch::jit::Module& module, const std::string& method_name) {
@@ -136,11 +144,7 @@ PYBIND11_MODULE(_C, m) {
     .def(py::init<>())
     .def_readwrite("min", &InputRange::min)
     .def_readwrite("opt", &InputRange::opt)
-    .def_readwrite("max", &InputRange::max)
-    .def("_to_internal_input_range", &InputRange::toInternalInputRange);
-
-  //py::class_<core::conversion::InputRange>(m, "_InternalInputRange")
-  //    .def(py::init<>());
+    .def_readwrite("max", &InputRange::max);
 
   py::enum_<DataType>(m, "dtype")
     .value("float",   DataType::kFloat)