trt_importer.py

# Copyright 2020, Visual Computing Group at HTW. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import abc
import time
import tensorrt as trt
import tensorflow as tf

from viscoop.tensorrt.trt_builder import TRTNetworkBuilder
from viscoop.tensorrt.tf_parser import TFProtobufParser
from viscoop.tensorrt.trt_inference import TRTInference

from tensorflow.python.platform import gfile


class TRTImporter(object):
    """

    Similar to
    https://github.com/onnx/onnx-tensorrt/blob/6.0-full-dims/ModelImporter.cpp#L457
    https://github.com/onnx/onnx-tensorrt/blob/ba53ee59da21af3096e38721327c74ec689f0f07/ModelImporter.cpp#L455

    """

    __metaclass__ = abc.ABCMeta

    def __init__(self, trt_logger_severity=trt.Logger.INFO):
        """

        :param trt_logger_severity:
        """

        self._trt_logger_severity = trt_logger_severity
        self._TRT_LOGGER = trt.Logger(trt_logger_severity)
        trt.init_libnvinfer_plugins(self._TRT_LOGGER, "")

    def from_tensorflow_pb_file(self, pb_file, input_node_names, input_node_shapes, output_node_names):
        """
        Convert the graph in the pb file to a trt network

        :param pb_file:
        :param input_node_names:
        :param input_node_shapes:
        :param output_node_names:
        :return: TRT Network
        """

        # read pb file
        graph_def = tf.GraphDef()
        with gfile.FastGFile(pb_file, 'rb') as f:
            graph_def.ParseFromString(f.read())
        return self.from_tensorflow_graph_def(graph_def, input_node_names, input_node_shapes, output_node_names)

    def from_tensorflow_graph_def(self, frozen_graph_def, input_node_names, input_node_shapes, output_node_names):

        # setup TRT network builder
        with trt.Builder(self._TRT_LOGGER) as builder:
            network = builder.create_network()
        network_builder = TRTNetworkBuilder(network)
        network_builder.register_inputs(input_node_names, input_node_shapes)
        network_builder.register_outputs(output_node_names)

        # parse Tensorflow PB file and build TRT network
        parser = TFProtobufParser(frozen_graph_def, network_builder)
        if parser.parse():
            return network_builder.network

        # Something went wrong
        return None

    def optimize_network(self, network, max_batch_size=1, fp16_mode=False, max_workspace_size=4294967296, fast_pass=False):
        """
        Optimize a TRT network and create a TRT engine out of it

        :param network:
        :param max_batch_size:
        :param fp16_mode:
        :param max_workspace_size:
        :param fast_pass:
        :return: TRT Engine
        """

        with trt.Builder(self._TRT_LOGGER) as builder:
            builder.max_batch_size = max_batch_size
            builder.max_workspace_size = max_workspace_size
            builder.fp16_mode = fp16_mode
            builder.min_find_iterations = 1 if fast_pass else 5
            builder.average_find_iterations = 1 if fast_pass else 5
            print("max_workspace_size", builder.max_workspace_size)
            print("average_find_iterations", builder.average_find_iterations)
            print("min_find_iterations", builder.min_find_iterations)

            # build and serialize the network
            st = time.time()
            engine = builder.build_cuda_engine(network)
            print('finished building an engine in {0:.2f} [sec]'.format((time.time() - st)))
            print("num_layers", engine.num_layers)
            print("device_memory_size", engine.device_memory_size)
            return engine.serialize()

    def inference_engine(self, engine):
        """
        Setup and return an inference context

        :param engine:
        :return: inference context
        """
        return TRTInference(engine, self._trt_logger_severity)

    @staticmethod
    def store_engine(engine, engine_file):
        """
        Store a TRT engine to the hard drive

        :param engine:
        :param engine_file:
        :return: bool
        """

        with open(engine_file, "wb") as f:
            num_chars = f.write(engine)
        print("stored TensorRT engine", engine_file)
        return num_chars > 1

    @staticmethod
    def load_engine(engine_file):
        """
        Load a TRT engine from the hard drive

        :param engine_file:
        :return: TRT Engine
        """

        # read from file
        with open(engine_file, "rb") as f:
            return f.read()