[TensorRT] Builder 에서의 UFF vs ONNX

환경

Ubuntu 16.04

Python 3.5

TensorRT 7.1.3.4

TensorFlow 1.13.1 및 1.15.1 

 

 

 

TensorFlow 프레임워크를 이용하여 .pb 를 만든 뒤 

convert-to-uff 유틸리티를 통해 .uff 파일로 변환하고나서 

.uff 파일을 tensorrt engine 으로 변환하는데

 

input order 와 input value 그리고 input dimension 을 올바르게 설정했는데도 불구하고 자꾸 아래와 같은 에러가 떠서 

 

[TensorRT] ERROR: UffParser: Parser error: tower_0/resnet_v1_50/conv1/BatchNorm/FusedBatchNorm: The input to the Scale Layer is required to have a minimum of 3 dimensions.

 

pb -> uff 파일 자체가 예전 TensorFlow 1.13.1 구버전을 이용해 생성한 것이라서 TensorRT 단에서 UFF 파일 못읽어 오는거라 생각해서 uff 파일을 TensorFlow 1.15.1 을 통해 pb 를 다시 만들고 TensorRT 7.1.3.4 단에서 다시 convert-to-uff 유틸을 통해 uff 파일을 만들고 다시 시도 하였는데 역시나 위와 같은 에러가 나서 찬찬히 살펴보니...

 

uff 파일의 network 를 아예 못읽어 오는 것이였다. uff 파일을 visualization 하여 제대로 생성되었음을 확인하였는데도 불구하고 network 를 못읽어온다는게 이상해서 다시 한번 찬찬히 살펴보니

 

uff 모델은 explicit batch 와 상관 없는 아이 였다. 

 

 

[ONNX 모델을 engine 으로 생성하는 코드 ]

 

# The Onnx path is used for Onnx models.
def build_engine_onnx(model_file):
    with trt.Builder(TRT_LOGGER) as builder, builder.create_network(common.EXPLICIT_BATCH) as network, trt.OnnxParser(network, TRT_LOGGER) as parser:
        builder.max_workspace_size = common.GiB(1)
        # Load the Onnx model and parse it in order to populate the TensorRT network.
        with open(model_file, 'rb') as model:
            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
        return builder.build_cuda_engine(network)

 

 

[UFF 모델을 engine 으로 생성하는 코드]

 

# The UFF path is used for TensorFlow models. You can convert a frozen TensorFlow graph to UFF using the included convert-to-uff utility.
def build_engine_uff(model_file):
    # You can set the logger severity higher to suppress messages (or lower to display more messages).
    with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.UffParser() as parser:
        # Workspace size is the maximum amount of memory available to the builder while building an engine.
        # It should generally be set as high as possible.
        builder.max_workspace_size = common.GiB(1)
        # We need to manually register the input and output nodes for UFF.
        parser.register_input(ModelData.INPUT_NAME, ModelData.INPUT_SHAPE)
        parser.register_output(ModelData.OUTPUT_NAME)
        # Load the UFF model and parse it in order to populate the TensorRT network.
        parser.parse(model_file, network)
        # Build and return an engine.
        return builder.build_cuda_engine(network)

 

onnx 모델을 engine 으로 생성하는 코드를 uff 모델용으로 바꾸면서 생긴 실수이다. 

 

이외에도 맨 처음에 언급했던 문제는 input order 가 잘못되었을 때 나타날 수 있다고 한다. 

 

 

 

참고자료 1 : forums.developer.nvidia.com/t/uff-parser-errors/62641

UFF parser errors