1. ONNX Concepts

OpenJR·2024년 4월 26일

저는 자율주행을 위한 딥러닝을 다루다 보니,

경량화를 위한 모델 아키텍처링이나
이 구조가 실제 제품에 올라갈 수 있을까?
라는 관점에 더 주목하게 됩니다.

그러다보니 onnx, tensorrt와 같이 Deployment를 위한 라이브러리에 관심이 가게 되어 그 중 ONNX에 대해 정리해보고자 합니다.

ONNX 공식문서만을 참고하였으며 필요한 내용으로 요약하였음.

ONNX Concepts

엔지니어들과 같이 일하다 보면 "onnx 그래프를 그리면~" 과 같이 그래프라는 용어가 자주 들장하게 되어 항상 들을 때 마다 솔직히 궁금했습니다. 대체 뭔말이지? 딥러닝 모델에 왠 그래프?
ONNX 공식문서에서는 그래프를 그린다는 것을 아래와 같이 표현하더군요.

"Building an ONNX graph means implementing a function with the ONNX language or more precisely the ONNX Operators."

네 그렇습니다. 단순히 Pytorch나 다른 프레임워크로 학습으로 생성한 모델을 ONNX Operator 함수로 구현한다는 말입니다.
부연 설명하자면, 학습한 모델의 Conv, MLP, ReLU 등의 Operation을 ONNX Operation으로 대체? 변환? 한다 정도인 것 같습니다.
근데 왜 이 과정을 그래프를 빌드한다 라고 표현했을까요?

1. Input, Output, Node, Initializer, Attributes

그 이유는 ONNX Operation (앞으로는 연산자라고도 표현하겠습니다.) 각각을 하나의 Node처럼 여기기 때문입니다.
아래의 Matmul과 Add란 연산자를 Input과 Ouput이 존재하는 (당연한...) Node로 표현하며, 이 노드들은 다음의 속성 값들이 꼭 필요합니다.

Input: float[M,K] x, float[K,N] a, float[N] c
Output: float[M, N] y

r = onnx.MatMul(x, a)
y = onnx.Add(r, c)

type:이 노드가 어떤 ONNX 연산자인지
Input: 노드의 입력
Output: 노드의 출력
Attribute: 특정 ONNX 연산자는 상수 값을 가지고 있으여 이 상수가 Attribute에 해당 (예시) BatchNorm)

2. Serialization with protobuf

ONNX로 배포할 때는, 모델을 Protobuf를 사용해 직렬화를 해서 압축(최적화)한다고 합니다.
"It aims at optimizing the model size as much as possible."

파이썬에서 onnx 사용하면 가끔 씩 다른 파이썬 라이브러리들이란 protobuf 관련된 디펜던시 문제가 생기던데, 그 이유가 이것 때문인 것 같네요.

3. List of available operators and domains

모델을 학습할 때, 어떤 Operation으로 학습했는지에 따라 ONNX 변환이 안될 때가 있습니다. (보통은 엄청 최신 모델이 아닌 이상 되긴 함/또는 배포할 디바이스의 onnx 버전이 너무 낮으면 또 안되기도 함)
그래서 배포를 염두에 두고 모델 개발할때는 항상 이 연산자가 onnx 변환이 가능한지 확인을 해야하는데요.
이 ONNX Operators 페이지에서 확인이 가능합니다.
신기한 거는 이 변환 가능한 연산자 집합을 무슨 도메인으로 식별하도록 나눠놨다고 하는데 별 내용은 제가 이해하기론 딱히 없는거 같애요.

ai.onnx: DL용 (MLP, Conv, ReLU 등등)
ai.onnx.ml: ML용 (TreeEnsemble Regressor 등등), 거의 신경안쓸 듯
opset이라는 개념도 있는데 이건 뒤에서 같이 알아보도록 하죠.

4. Supported Types

ONNX는 Tensor만 내부 연산이 가능하다고 합니다. 당연히 입력도 Tensor 형태여야겠죠?
굳이 텐서를 정의하자면,

type: int인지 float인지
shape: tensor의 차원
contiguous array: 텐서는 연속 배열이여야 한다.
라고 합니다.

4.1. Element Type

기본적으로 타입은 딥러닝 국룰 기본값인 float32라고 합니다. 추가적으로 아래의 타입들도 지원한다고 합니다.

 0: onnx.TensorProto.UNDEFINED
 1: onnx.TensorProto.FLOAT
 2: onnx.TensorProto.UINT8
 3: onnx.TensorProto.INT8
 4: onnx.TensorProto.UINT16
 5: onnx.TensorProto.INT16
 6: onnx.TensorProto.INT32
 7: onnx.TensorProto.INT64
 8: onnx.TensorProto.STRING
 9: onnx.TensorProto.BOOL
10: onnx.TensorProto.FLOAT16
11: onnx.TensorProto.DOUBLE
12: onnx.TensorProto.UINT32
13: onnx.TensorProto.UINT64
14: onnx.TensorProto.COMPLEX64
15: onnx.TensorProto.COMPLEX128
16: onnx.TensorProto.BFLOAT16
17: onnx.TensorProto.FLOAT8E4M3FN
18: onnx.TensorProto.FLOAT8E4M3FNUZ
19: onnx.TensorProto.FLOAT8E5M2
20: onnx.TensorProto.FLOAT8E5M2FNUZ
21: onnx.TensorProto.UINT4
22: onnx.TensorProto.INT4

공식 문서에서는 numpy와 onnx간 매핑함수를 지원하여 아래와 같이 매핑 값을 알 수 있다고 하는데요, 아마 pytorch나 tensorflow에서 이 값들을 땡겨와 자기네들 모델을 onnx로 변환할 때 사용하는 것 같네요.

import re
from onnx import TensorProto

reg = re.compile('^[0-9A-Z_]+$')

values = {}
for att in sorted(dir(TensorProto)):
    if att in {'DESCRIPTOR'}:
        continue
    if reg.match(att):
        values[getattr(TensorProto, att)] = att
for i, att in sorted(values.items()):
    si = str(i)
    if len(si) == 1:
        si = " " + si
    print("%s: onnx.TensorProto.%s" % (si, att))

4.2. Sparse Tensor & Other types

Sparse Tensor도 지원한다곤 하는데 2차원 까지 밖에 안하고, SequenceProto, MapProto 같은 다른 타입도 있다고 합니다.

5. What is an opset version?

이제 Opset 이야기를 해보려 합니다. 저는 거의 99.9% 모델 학습 시 Pytorch를 사용하는데요, Pytorch 모델을 ONNX 파일로 변환하기 위해 아래와 같이 코드를 작성합니다.

torch.onnx.export(
    model,
    dummy_input,
    onnx_model_path,
    verbose=True,
    input_names=model_inputs,
    output_names=["features"],
    opset_version=14  # ONNX opset version (adjust as needed)
    )

다른 인자들을 신경쓰지 마시고, opset_version만 보시면 됩니다.
위 처럼 거의 onnx로 변환할 때, opset_version을 명시해서 적어 줍니다.
opset이 대체 뭐냐면, 위에서 ONNX Operators를 설명드릴 때 나왔던 onnx로 변환 가능한 연산자들의 버전을 의미합니다.
opset이 증가함에 따라 연산자들은 업데이트 되거나, 새로운 연산자가 추가 됩니다.
이 때, 해당 연산자가 몇 Opset 버전에서 업데이트 되거타 추가되었는지 확인하기위해 opset이 필요합니다.

6. Subgraphs, tests and loops

Pytorch 모델의 Forward 내부에서 for문이나 if문이 있다면 되기는 하는데 매우 느려져서 피하는게 좋다고 합니다. For와 If문을 지양합시다!

7. Extensibility

ONNX Operators에서 ONNX에 사용가능한 연산자들을 표준으로 정의했지만, 새로운 연산자를 커스텀해서 사용할 수는 있다고 합니다.
보통 근데 ONNX를 TensorRT로 변환하기 때문에, 이러한 경우 적용 안될 것 같긴 합니다만, ONNX Runtime만 사용할 때에는 유용할 것으로 보입니다.

8. Shape (and Type) Inference

텐서의 Shape을 아는 것이 ONNX Graph를 실행하는데에 꼭 필수 요소는 아니라고 합니다.
하지만 알면 메모리의 위치를 기억해서 추론하는데 더 속도가 따르다고 하네요.

9. Tools

마지막입니다. Netron을 추천한다고 합니다. 사용하면 아래처럼 그래프를 시각화해서 이해하는데 도움이 된다고 합니다.

와 이건 저도 몰랐는데, zetane라는 툴도 있다고 합니다. 이 툴을 사용하면 onnx를 실행할 때 중간 결과도 아래와 같이 볼 수 있다고 하네요.