From 96c13f767ba896b43eadfd799f6427d6ebd9b2ac Mon Sep 17 00:00:00 2001
From: Heyufan <yf.he@caicloud.io>
Date: Wed, 12 Aug 2020 18:00:52 +0800
Subject: [PATCH 1/2] feat: support tensorrt extract

---
 build/extract/tensorrt/Dockerfile             | 28 ++++++++++
 scripts/extract/extractor/__init__.py         |  4 +-
 scripts/extract/extractor/extract_tensorrt.py | 51 +++++++++++++++++++
 3 files changed, 82 insertions(+), 1 deletion(-)
 create mode 100644 build/extract/tensorrt/Dockerfile
 create mode 100644 scripts/extract/extractor/extract_tensorrt.py

diff --git a/build/extract/tensorrt/Dockerfile b/build/extract/tensorrt/Dockerfile
new file mode 100644
index 00000000..448c0aaa
--- /dev/null
+++ b/build/extract/tensorrt/Dockerfile
@@ -0,0 +1,28 @@
+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
+
diff --git a/scripts/extract/extractor/__init__.py b/scripts/extract/extractor/__init__.py
index 6e34f39c..d95da7c7 100644
--- a/scripts/extract/extractor/__init__.py
+++ b/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]')
\ No newline at end of file
+    raise ImportError('module must be in one of [onnx, caffe, netdef, graphdef, h5, mxnetparams, savedmodel, torchscript, tensorrt]')
\ No newline at end of file
diff --git a/scripts/extract/extractor/extract_tensorrt.py b/scripts/extract/extractor/extract_tensorrt.py
new file mode 100644
index 00000000..56be46f6
--- /dev/null
+++ b/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())
\ No newline at end of file

From 4780a604f8519f3a524b7b27ee82b4edec93f4c7 Mon Sep 17 00:00:00 2001
From: Yufan He <30286093+judgeeeeee@users.noreply.github.com>
Date: Wed, 12 Aug 2020 19:02:03 +0800
Subject: [PATCH 2/2] docs: update build/extract/tensorrt/Dockerfile

Co-authored-by: Ce Gao <ce.gao@outlook.com>
---
 build/extract/tensorrt/Dockerfile | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/build/extract/tensorrt/Dockerfile b/build/extract/tensorrt/Dockerfile
index 448c0aaa..71ba3564 100644
--- a/build/extract/tensorrt/Dockerfile
+++ b/build/extract/tensorrt/Dockerfile
@@ -21,8 +21,7 @@ RUN apt update -y \
 RUN ln -sf /usr/share/zoneinfo/Asia/Shanghai /etc/localtime
 
 COPY scripts/shell/run.sh /scripts/run.sh
-COPY scripts/extract  /scripts
+COPY scripts/extract /scripts
 WORKDIR /scripts
 
 ENV EXTRACTOR=tensorrt
-