diff --git a/build/extract/tensorrt/Dockerfile b/build/extract/tensorrt/Dockerfile new file mode 100644 index 00000000..71ba3564 --- /dev/null +++ b/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 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