feat: support tensorrt extract

build/extract/tensorrt/Dockerfile

@@ -0,0 +1,27 @@
+FROM nvcr.io/nvidia/tensorrt:19.10-py3
+ENV LC_ALL="C.UTF-8" \
+  LANG="C.UTF-8"
+
+ARG ormb_version=v0.0.3
+ARG ormb_tar_filename=ormb_0.0.3_Linux_x86_64.tar.gz
+
+RUN apt update -y \
+    && apt install wget -y \
+    && pip config set global.index-url https://mirrors.aliyun.com/pypi/simple/ \
+    && pip install --no-cache-dir \
+                grpcio==1.26.0\
+                grpcio-tools\
+                future\
+                pyyaml\
+    && wget https://github.com/kleveross/ormb/releases/download/$ormb_version/$ormb_tar_filename \
+    && tar -xvf $ormb_tar_filename -C /usr/local/bin \
+    && rm -rf $ormb_tar_filename
+
+#Set timezone
+RUN ln -sf /usr/share/zoneinfo/Asia/Shanghai /etc/localtime
+
+COPY scripts/shell/run.sh /scripts/run.sh
+COPY scripts/extract /scripts
+WORKDIR /scripts
+
+ENV EXTRACTOR=tensorrt

scripts/extract/extractor/__init__.py

@@ -19,5 +19,7 @@
     from .extract_torchscript import TorchscriptExtractor as Extractor
 elif module=='pmml':
     from .extract_pmml import PMMLExtractor as Extractor
+elif module=='tensorrt':
+    from .extract_tensorrt import TensorrtExtractor as Extractor
 else:
-    raise ImportError('module must be in one of [onnx, caffe, netdef, graphdef, h5, mxnetparams, savedmodel, torchscript]')
+    raise ImportError('module must be in one of [onnx, caffe, netdef, graphdef, h5, mxnetparams, savedmodel, torchscript, tensorrt]')

scripts/extract/extractor/extract_tensorrt.py

@@ -0,0 +1,51 @@
+import os
+import json
+import collections
+
+import tensorrt as trt
+
+from .base_extract import BaseExtrctor
+
+MODEL_TYPE = 'TensorRT'
+MODEL_EXTENSION1 = '.engine'
+MODEL_EXTENSION2 = '.plan'
+
+
+class TensorrtExtractor(BaseExtrctor):
+    def _extract_inputs(self):
+        engine = self.engine
+        inputs = []
+        for binding in engine:
+            if engine.binding_is_input(binding):
+                origin_inputs = {}
+                origin_inputs['name'] = binding
+                origin_inputs['dtype'] = trt.nptype(
+                    engine.get_binding_dtype(binding)).__name__
+                origin_inputs['size'] = engine.get_binding_shape(binding)
+                inputs.append(origin_inputs)
+        return inputs
+
+    def _extract_outputs(self):
+        engine = self.engine
+        outputs = []
+        for binding in engine:
+            if not engine.binding_is_input(binding):
+                origin_outputs = {}
+                origin_outputs['name'] = binding
+                origin_outputs['dtype'] = trt.nptype(
+                    engine.get_binding_dtype(binding)).__name__
+                origin_outputs['size'] = engine.get_binding_shape(binding)
+                outputs.append(origin_outputs)
+        return outputs
+
+    def _extract_ops(self):
+        return {'layers': self.engine.num_layers}
+
+    def _load_model(self):
+        try:
+            model_path = self._find_with_extension(MODEL_EXTENSION1)
+        except:
+            model_path = self._find_with_extension(MODEL_EXTENSION2)
+        TRT_LOGGER = trt.Logger()
+        with open(model_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
+            self.engine = runtime.deserialize_cuda_engine(f.read())