уменьшение скорости при преобразовании модели Tensorflow в TensorRT с использованием синтаксического анализатора ONNX - PullRequest
Новая
       23

уменьшение скорости при преобразовании модели Tensorflow в TensorRT с использованием синтаксического анализатора ONNX

0 голосов
/ 16 февраля 2020

Я обучил пользовательской модели CNN, которая состоит из нескольких общих слоев, таких как conv / maxpool / f c, ... Я хочу преобразовать модель TensorFlow в чистый TensorRT, используя парсеры ONNX и UFF, но когда я использую парсер ONNX для чистого двигателя TensorRT я теряю скорость, но с парсером UFF скорость модели увеличивается, почему? Я выполнил следующие шаги: 1- 

python3 -m tf2onnx.convert --graphdef my_model.pb --output my_model.onnx \
    --fold_const --opset 10 \
    --inputs the_input:0 \
    --outputs the_output:0

2- 

def build_engine(onnx_file_path, engine_file_path, verbose=False):
    """Takes an ONNX file and creates a TensorRT engine."""
    TRT_LOGGER = trt.Logger(trt.Logger.VERBOSE) if verbose else trt.Logger()
    with trt.Builder(TRT_LOGGER) as builder, builder.create_network(*EXPLICIT_BATCH) as network, trt.OnnxParser(network, TRT_LOGGER) as parser:
        builder.max_workspace_size = 1 << 28
        builder.max_batch_size = 1
        builder.fp16_mode = True
        #builder.strict_type_constraints = True

        # Parse model file
        if not os.path.exists(onnx_file_path):
            print('ONNX file {} not found, please run yolov3_to_onnx.py first to generate it.'.format(onnx_file_path))
            exit(0)
        print('Loading ONNX file from path {}...'.format(onnx_file_path))
        with open(onnx_file_path, 'rb') as model:
            print('Beginning ONNX file parsing')
            if not parser.parse(model.read()):
                print('ERROR: Failed to parse the ONNX file.')
                for error in range(parser.num_errors):
                    print(parser.get_error(error))
                return None
        if trt.__version__[0] >= '7':
            # The actual yolov3.onnx is generated with batch size 64.
            # Reshape input to batch size 1
            shape = list(network.get_input(0).shape)
            shape[0] = 1
            network.get_input(0).shape = shape
        print('Completed parsing of ONNX file')

        print('Building an engine; this may take a while...')
        engine = builder.build_cuda_engine(network)
        print('Completed creating engine')
        with open(engine_file_path, 'wb') as f:
            f.write(engine.serialize())
        return engine
...