[TensorRT] ONNX 및 TRT에서 Group Normalization 사용하기 (+ Instance Normalization 이슈)

 

PyTorch 1.6.0 / ONNX opset 10, 11 / TensorRT 7.2.2.3 환경에서

파이토치의 어떤 모델을 ONNX 모델로 변환하고, ONNX 모델을 TensorRT 모델(+ Dynamic Shape)로 변환한 뒤

Python, C++ 각 환경에서 사용하고자 할 때 아래와 같이 Engine Serialization를 못하는 문제가 발생하였었다. 

 

결론부터 말하자면 TensorRT에서는 Group Normalization을 지원하지 않는다. 

 

PyTorch에서 Group Normalization을 사용하게 되면

ONNX와 TensorRT에서 Instance Normalization 노드가 생기게 되는데

이는 TensorRT 에서 지원하지 않는다. Custom Plugin 을 생성해줘야한다.

 

ONNX 및 TRT에서 Group Normalization 사용하는 방법은 간단히 말하자면 아래와 같다. 

 

1. ONNX-GS를 이용하여 ONNX Graph 변환 

 

원래 PyTorch 1.4 기준 Group Norm 레이어는 ONNX(opset 11)에서 다음과 그림과 같이

Reshape, Shape, Unsqueeze, Mul, Add, Instance Normalization 레이어로 구성되어있다. 

이를 ONNX-GS를 이용하여 Custom Layer를 생성하면 단일 레이어로 축소되고, CUDA 커널에서 사용할 수 있다. 

출처 : NVIDIA

노드를 대치하는 코드는 아래와 같다.

 

import onnx_graphsurgeon as gs
 
 
def process_groupnorm_nodes(graph):
   """
   Gather the instance normalization nodes and the rest of the subgraph
   and convert into a single GN node
   """
   instancenorms = [node for node in graph.nodes if node.op == "InstanceNormalization"]
   for node in instancenorms:
        convert_to_groupnorm(node, graph)
 
    return graph
 
 
def convert_to_groupnorm(instancenorm, graph):
   """
    Convert the PyTorch-exported GroupNorm subgraph to the subgraph below
    Conv
      |
    GroupNorm
      |
    ReLU
    Attributes:
        instancenorm: Instance Normalization node in the graph.
        graph: Input graph object
    """
 
   attrs = retrieve_attrs(instancenorm)
   groupnorm = gs.Node(op="GroupNormalizationPlugin", attrs=attrs)
   graph.nodes.append(groupnorm)
  
   # The plugin receives an input from the Conv node, and output to the ReLU node
   # o() corresponds to the node output and i() corresponds to node input.
   # Output of Conv
   conv_output_tensor = instancenorm.i().inputs[0]
   # Output of Add.
   relu_input_tensor = instancenorm.o().o().o().outputs[0]
 
   # Reconnect inputs/outputs to the groupnorm plugin
   conv_output_tensor.outputs[0] = groupnorm
   relu_input_tensor.inputs[0] = groupnorm
 
   # Add scale and bias tensors from unsqueeze op as input to groupnorm plugin
   groupnorm.inputs.append(instancenorm.o().o().i(1).inputs[0])
   groupnorm.inputs.append(instancenorm.o().o().o().i(1).inputs[0])

 

이렇게 복잡했던 노드들을 GroupNormalizationPlugin 노드로 대체하게 되면 아래와 같이 깔끔하게 모델이 생성된다. 

 

 

노드를 대치하여 그래프를 생성하고, ONNX 모델을 저장하는 코드는 아래와 같다. 

 

import torch
import onnx
from monodepth.models.networks.PackNet01 import PackNet01
 
def post_process_packnet(model_file, opset=11):
 
   """
    Use ONNX-Graphsurgeon to replace upsample and instance normalization nodes. For more information, see the post_processing.py script
    Args:
        model_file : Path to ONNX file
    """

    # Load the PackNet graph
    graph = gs.import_onnx(onnx.load(model_file))

    if opset==11:
        graph = process_pad_nodes(graph)

    # Replace the subgraph of upsample with a single node with input and scale factor.
    graph = process_upsample_nodes(graph, opset)

    # Convert the GN subgraph into a single plugin node.
    graph = process_groupnorm_nodes(graph)

    # Remove unused nodes, and topologically sort the graph.
    graph.cleanup().toposort()

    # Export the ONNX graph from graphsurgeon
    onnx.save_model(gs.export_onnx(graph), model_file)

    print("Saving the ONNX model to {}".format(model_file))

 

 

 

2. TRT에서 Group Normalization Plugin 구현

 

TensorRT에서는 GN 레이어를 지원하지 않기 때문에 Plugin을 구현해주어야 한다. 

TRT GitHub에 여러 구현들이 포함되어 있다. https://github.com/NVIDIA/TensorRT/tree/master/plugin

구현에는 레이어의 핵심 구현이 포함 되어있는 GroupNormalizationPlugin 클래스와

플러그인을 생성하고 매개변수를 설정하는 GroupNormalizationPluginCreator 클래스가 존재한다. 

 

참고로 TRT 7.1 버전 부터는 플러그인을 레지스트리에 정적으로 등록하여 사용할 수 있다고 한다. 

 

REGISTER_TENSORRT_PLUGIN(GroupNormalizationPluginCreator)

 

플러그인을 구현한 후, 컴파일을 위해 CMakeLists.txt 및 plugin/InferPlugin.cpp에 플러그인 종속성을 추가해준다. 

그 다음 README.md 에 따라 빌드를 진행하여 libnvinfer_plugin.so 를 생성한다.

이 라이브러리의 경로를 LD_LIBRARY_PATH에 추가해준다. 

 

그 다음 trtexec를 사용하여 엔진 생성 및 추론을 수행한다.

참고로 trtexec 에는 외부 플러그인 라이브러리를 로드하는 --plugins 옵션이 있다. 

 

 

자세한 사항은 NVIDIA 블로그에 상세히 설명되어 있다.

https://developer.nvidia.com/blog/estimating-depth-beyond-2d-using-custom-layers-on-tensorrt-and-onnx-models/

Estimating Depth with ONNX Models and Custom Layers Using NVIDIA TensorRT | NVIDIA Technical Blog

 

 

 

 

 

 

 

3. 겪었던 이슈 

 

처음에 plugin 추가 없이 엔진을 생성하고 사용했을 때 엔진 빌드 자체가 안되서 구글링 해보니

플러그인을 임의로 초기화 해주면 된다는 글을 보았었다.

그래서 엔진을 사용할 때 Plugin을 생성하지 않고 Plugin을 임의로 초기화 하여 Batch Size = 1 사이즈를 가지는 엔진을 생성하고 추론 했더니 결과는 정상적으로 나왔다. 명시적 배치사이즈에서는 제대로 작동한다고 한다.

그런데 이는 Output Node가 PyTorch와 ONNX에서 지정한 순서대로 나오지 않는다는 문제가 있었다. 별도의 재 정렬이 필요하다.

 

 

그리고 또한 이 과정에서 엔진을 생성 할 때

관련된 Plugin을 추가하여 TRT를 빌드하지 않고 기존의 trtexec을 사용하여 엔진을 생성하게 되면 Batch Size = 1에서는 엔진이 생성 되지만, Batch Size >= 2 이상의 사이즈를 가지는 엔진은 생성 자체가 안된다. 정상적으로 사용하려면 Instance Norm Plugin을 포함하여 TRT를 빌드하고, 그렇게 빌드된 trtexec를 사용하여 엔진을 생성해 사용해야한다. 단, TRT 7.2 이상의 버전이어야한다. 이전 버전들에서는 Instance Norm 노드가 포함된 Dynamic Shape을 가지는 모델은 변환이 안된다고 한다. 

 

또한 간혹 Shape 지정 시 Input Name을 잘못 지정하고 Batch Size >= 2이상의 엔진을 생성하면

임의로 배치사이즈 1의 엔진이 생성되고, 이를 추론을 하게 되면 잘못된 결과를 초래하게 된다. 

 

그리고 참고로 TensorRT 8.x 버전에서는 Instance Normalization 을 사용하였을 때 성능 손실이 존재(해결중)한다고 한다. 

 

따라서 필자는 추후 TRT 8.x를 사용할 수도 있는 등 여러 이슈 때문에

아쉽지만 PyTorch 에서 모델을 구현 할 때 Group Normalization을 사용하지 않기로 하였다. 

 

그래서 이러한 과정을 아래와 같이 설명한다. 

 

 

3.1 Instance Normalization의 Plugin을 임의로 초기화할 때 발생하는 이슈 

 

TensorRT에서 Group Normalization을 지원하지 않기 때문에 엔진 빌드 시 아래와 같은 에러를 마주하게 된다. 

 

[TensorRT] ERROR: INVALID_ARGUMENT: getPluginCreator could not find plugin InstanceNormalization_TRT version 1
[TensorRT] ERROR: safeDeserializationUtils.cpp (322) - Serialization Error in load: 0 (Cannot deserialize plugin since corresponding IPluginCreator not found in Plugin Registry)
[TensorRT] ERROR: INVALID_STATE: std::exception
[TensorRT] ERROR: INVALID_CONFIG: Deserialize the cuda engine failed.

 

그래서 아래와 같은 형식으로 플러그인을 초기화 해주면 실행이 제대로 되었었다.

단 Batch Size = 1 일 때만 제대로 동작하는 것 같다.

(내 생각에는 다른 이슈도 있을 듯 하다. 이는 임시 방편으로만 사용하길 바라며, 비추천이다.)

 

trt.init_libnvinfer_plugins(TRT_LOGGER, "")

 

참고자료 1 : https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#add_custom_layer_python
참고자료 2 : https://forums.developer.nvidia.com/t/crash-when-converting-onnx-reid-model-to-tensorrt/109345/12

 

 

 

3.2 기존 trtexec를 이용하여 Batchsize 2 이상으로 엔진을 생성할 경우에 발생하는 이슈

 

Engine 자체가 생성되지 않는다.

[E] [TRT] instanceNormalizationPlugin.cpp (283) - Cuda Error in configurePlugin: 2 (out of memory) terminate called after throwing an instance of 'nvinfer1::plugin::CudaError' what(): std::exception Aborted (core dumped)

 

참고로 dynamic shape 지정 시 name 을 잘못 지정하면 임의로 batch size 1 로 셋팅되어 변환된다.

Dynamic dimensions required for input: input_name, but no shapes were provided.
Automatically overriding shape to ~~~ 

이렇게 이름을 잘못 지정하게 되면 optShapes, maxShapes, minShapes 을 어떻게 지정하느냐에 따라 다르겠지만

이상한 결과 값을 내뱉는다. (Inference Deley, Output Feature Map 잘못된 값)

관련 이슈 : https://github.com/NVIDIA/TensorRT/issues/1111

 

 

 

아무튼 Instance Norm에 관련된 이슈를 찾아보니

TRT 7.2 이전 버전에서 Dynamic Shape 으로 Engine 을 생성할 때 변환이 안되는 이슈가 있었다. 

Instance Norm 을 사용할 때 느리다는 이슈도 있었다. 

 

관련 이슈 1 : InstanceNormalization does not support dynamic inputs ( https://github.com/onnx/onnx-tensorrt/issues/374 )

관련 이슈 2 : TensorRT InstanceNormalization is too slow ( https://github.com/NVIDIA/TensorRT/issues/595 )

 

이는 아래와 같이 TRT 7.2에서 수정되었다고 한다.

관련 이슈 : https://github.com/NVIDIA/TensorRT/issues/1222

 

수정된 TRT 7.2 내용은 아래와 같다. 

https://github.com/NVIDIA/TensorRT/blob/21.04/plugin/instanceNormalizationPlugin/instanceNormalizationPlugin.cpp

GitHub - NVIDIA/TensorRT: TensorRT is a C++ library for high performance inference on NVIDIA GPUs and deep learning accelerators

관련 이슈 : https://github.com/NVIDIA/TensorRT/issues/1183#issuecomment-828105388

 

 

 

3.3 TRT 8.x에서 발생하는 이슈 

 

또한 참고로 TensorRT Release-notes 에 따르면

TensorRT 8.x 에서 Instance Norm plugin 을 사용할 때 성능 손실이 있는 것으로 확인된다.

이 이슈는 NVIDIA 에서 해결중이라고 한다. 

 

There is an up to 25% performance drop for networks using the InstanceNorm plugin. This issue is being investigated.

 

 

 

+ 현재 날짜 기준(2022.02.23)으로 TensorRT 8.4.0 Early Access 버전에서는

8.x 버전에서 Instance Norm plugin 이슈 언급이 없어졌다.

그런데 Fixed Issues에는 명시되어 있지는 않다. 추후 사용할 때 확인이 필요할 듯 하다. 

https://docs.nvidia.com/deeplearning/tensorrt/release-notes/

Release Notes :: NVIDIA Deep Learning TensorRT Documentation

 

 

아무튼 ... 결론적으로 이러한 이슈가 잔존하는 것으로 보아

필자는 Instance Norm 레이어를 생성하게 되는 Pytorch의 구현을 바꾸어주는 것이 깔끔할 것 같다고 판단되었다.