[Deploy] ONNX: 다른 DNN 프레임워크 간 모델 호환 포멧(pytorch, tensorflow, TensorRT, ...)

💡 본 문서는 '[Perception] ONNX: 다른 DNN 프레임워크 간 모델 호환 포멧(pytorch, tensorflow, ...)'에 대해 정리해놓은 글입니다.

1. PyTorch, ONNX, TensorRT 비교

1) PyTorch 2.0

언젠가부터 Tensorflow를 앞지르고 가장 인기있는 Deep Learning Framework이 된 PyTorch가 2022년 12월에 새로운 버전을 공개했습니다. 이 글에서 알아야 할 중요한 부분은 PyTorch는 너무 Pythonic해서 문제라는 것입니다.

Python의 특징

그럼 Pythonic 하다는 게 뭐냐?

• 간결성 / 가독성: User friendly하게 작성하기 때문에 코드가 직관적이고 over-engineering을 막을 수 있습니다.
  • 반대로 말하면 컴퓨터가 알아듣기 어렵기 때문에 컴퓨터는 이것을 번역하느라 시간이 오래 걸립니다.
• 인터프리터 언어: 컴파일 언어와는 다르게 컴파일 과정 없이 매번 한줄한줄 실행합니다. 그래서 대표적인 컴파일 언어인 C에서는 뭐가 잘못 코딩되면 아예 실행을 할 수가 없지만 Python에서는 처음에 잘 실행되는 것 같다가도 중간에 에러를 발생시킵니다.
  • 미리 컴파일 하지 않기 때문에 컴퓨터는 매번 코드를 따라 새로운 모험을 하는 겁니다.. 실행 시간이 느리겠죠

이 밖에도 Python의 특징은 많이 있지만 아무튼 사용하기 쉽지만 그만큼 느리다라는 특징만 아시면 됩니다.

Pytorch 2.0의 특징

그래서 PyTorch 2.0에서는 torch.compile을 포함한 몇가지 기능들을 추가하여 더욱 빠르게 Inference 할 수 있게 업데이트가 되었습니다. 이걸 이해하려면 JIT(JustInTime)을 알고 있어야 하는데, 요약하자면 미리 컴파일 하여 그 뒤로는 빠르게 추론할 수 있음이 핵심입니다. 자세한 내용은 torch.compile tutorial을 참고해주세요.

2) ONNX

이미지 출처:  https://microsoft.github.io/ai-at-edge/docs/onnx/

ONNX는 "Open Neural Network Exchange"의 약어로, 오픈 소스 프로젝트입니다. ONNX는 인공지능 AI 모델을 표준 형식으로 표현하고 서로 다른 딥러닝 프레임워크 간에 모델을 변환하고 공유할 수 있게 해줍니다.

ONNX는 딥러닝 모델의 구조와 가중치를 표현하는 중립적인 형식입니다. 이 형식은 다양한 딥러닝 프레임워크(TensorFlow, PyTorch, Keras)로부터 모델을 내보내고, 다른 프레임워크에서 가져와서 실행하거나 변환할 수 있습니다. ONNX 형식을 사용하면 개발자들은 다양한 프레임워크를 유연하게 조합하고 모델을 재사용하고 다양한 언어, 플랫폼 및 디바이스에서 실행할 수 있습니다.

TensorRT

TensorRT도 ONNX와 마찬가지입니다. 하지만 이건 NVIDIA에서 만든 프레임워크로써, NVIDIA GPU에서 최적화 된 기술입니다. 그런데 현재까지는 대부분의 연구가 NVIDIA GPU를 사용하기 때문에 효과적인 방법이라고 할 수 있습니다.

이미지 출처:  https://developer.nvidia.com/ko-kr/blog/nvidia-tensorrt-inference-%EC%B5%9C%EC%A0%81%ED%99%94-%EB%B0%8F-%EA%B0%80%EC%86%8D%ED%99%94%EB%A5%BC-%EC%9C%84%ED%95%9C-nvidia%EC%9D%98-toolkit/

4) PyTorch 2.0 / ONNX / TensorRT 비교

그럼 이제 세가지 모델 형식을 비교하겠습니다. 참고로 실험 GPU는 NVIDIA RTX 3090Ti로, TensorRT에게 유리한 환경임을 참고하시기 바랍니다.

비교를 통해 얻을 수 있는 핵심 4가지는 아래와 같습니다.

• PyTorch 2.0은 batch size가 커질 수록 전작 대비 큰 성능 향상을 이룸
  • 또한 FP16은 Batch size가 클수록 빛을 발함
  • 이는 PyTorch 2.0이 학습 최적화에 초점을 맞췄기 때문
• ONNX Runtime은 Batch size가 작을 때 PyTorch 2.0보다 효과가 좋음
  • 이는 ONNX가 추론 최적화에 초점을 맞췄기 때문
• PyTorch Eager 기존모드나 PyTorch 2.0이나 Batch size가 작으면 별 차이가 없음
  • 이건 아마 GPU가 오버스펙이라 보유한 자원을 충분히 활용하지 못했기 때문
• NVIDIA GPU에서는 TensorRT가 최고
  • 당연한 것이지만 NVIDIA에서 최적화를 했기 때문에 Cache를 가장 효율적으로 활용하도록 설계됐으리라 추정

오늘은 위와 같이 PyTorch 2.0, ONNX, TensorRT 간의 추론 속도 비교를 해봤으며, 결론은 다음과 같습니다.

1. 학습할 때는 PyTorch 2.0 MixedPrecision 또는 FP16이 좋아보임
2. 추론할 때는 TensorRT가 좋아 보이긴 함. 하지만 딥러닝 모델을 배포할 때는 GPU를 쓰지 않을 수도, GPU가 NVIDIA 제품이 아닐 수도 있습니다. 잘 고려하셔야 합니다.
3. 다양한 환경이나 다른 프레임워크 간 호환성을 위해서는 모델을 ONNX 포멧으로 변환해야 합니다. 여기서 ONNX 포멧으로 변환하면 추론 최적화도 진행합니다.

2. ONNX

1) ONNX 란?

ONNX(Open Neural Network Exchange)는 Tensorflow, PyTorch 와 같은 서로 다른 DNN 프레임워크 환경에서 만들어진 모델들을 서로 호환해서 사용할 수 있도록 도와주는 공유 플랫폼입니다. 간략히 말해, 다양한 플랫폼 환경(Java, JS, C, C#, C++)에서 환경에 제약 없이 구현된 ‘ML 모델’을 호출하고 수행하여 수행 결과값을 반환받는 것을 의미합니다.

1. Framework Interoperability

위에서 언급했다시피 특정 환경에서 생성된 모델을 다른 환경으로 import하여 자유롭게 사용을 할 수 있다는 것은 ONNX의 최대 강점입니다. 예컨대, Tensorflow에서 빠르게 모델을 학습 시킨 뒤에 이를 모바일로 옮겨서 사용을 하는 등 여러가지 방식으로 활용 가능합니다.

2. Shared Optimization

HW vendor(가속기와 같은 HW 제조업체)의 입장에서 ONNX와 같은 프레임워크 간 공유되는 포맷이 존재하면, 하드웨어 설계시 ONNX representation을 기준으로 최적화를 하면 되기 때문에 효율적입니다.

마치 JSON 포맷이 정보 표현을 위해서 여러 개발자들 사이에서 합의되어 사용되듯, ONNX라는 합의된 DNN 모델 포맷이 존재한다고 생각하면 됩니다. ONNX 사용과 관련하여 보다 자세한 튜토리얼이 필요하다면 다음 페이지를 참고하면 됩니다.

2) ONNX Runtime 란?

ONNX 모델을 실행하기 위한 엔진입니다. ONNX 모델을 실행하기 위해 ONNX 런타임은 빠른 추론을 위한 최적화된 커널을 사용합니다. 또한, ONNX Runtime은 CPU, GPU 및 딥러닝 가속기(DNNL, NNAPI, OpenVINO)를 지원합니다. 따라서, ONNX Runtime은 ONNX 모델을 실행하기 위한 최적화된 런타임 환경을 제공합니다.

+ ORT 형식

• ‘축소된 크기’의 ONNX Runtime 빌드에서 지원하는 형식입니다.
• onnx 형태에서 ort 형태로 변경을 함으로써 축소된 크기 빌드는 모바일 및 웹 애플리케이션과 같이 크기가 제한된 환경에서 사용하기에 더 적합합니다.

3) ONNX 사용 사례

ONNX를 이용하여서 아래와 같이 사용할 수 있습니다.

• 1. 다양한 ML 모델에 대한 추론 성능 향상에 사용됩니다.
• 2. 다른 하드웨어 및 운영 체제에서 실행에 사용됩니다.
• 3. Python으로 훈련하되 C#/C++/Java 앱에 배포에 사용됩니다.
• 4. 다양한 프레임 워크에서 생성된 모델로 추론 훈련 및 수행에 사용됩니다.

3. ONNX 사용 과정

위 그림은, PyTorch 모델을 ONNX 그래프로 export 하는 전체 과정을 도식화한 것입니다.

진행 과정은 아래 순서와 같습니다.

• 첫 번째
  • PyTorch 모델과 Sample input 을 인자로 하여, torch.onnx.export 함수를 호출합니다.
  • PyTorch 의 JIT 컴파일러를 통해서, Trace 혹은 Script 를 생성합니다.
    • Trace 와 Script 는 그 생성 방식과 representation 에 차이가 있습니다. (추후 포스팅)
  • PyTorch 모델의 forward propagation 시에 호출되는, 함수 및 연산들에 대한 최적화된 그래프인 Torch IR 을 만듭니다.
    • Trace 나 Script 는, PyTorch 의 nn.Module 을 상속하는 모델의,
    • forward 함수에서 실행되는 코드들에 대한 IR(Intermediate Representation)을 담고 있습니다.
• 두 번째
  • 생성된 trace / script (Torch IR)는, ONNX Exporter 를 통해서 ONNX IR 로 변환되고,
  • 여기에서 한 번 더 Graph Optimization 이 이루어집니다.
• 세 번째
  • 최종적으로 생성된 ONNX 그래프는 .onnx 포맷으로 저장됩니다.

참고

• [Blog] ONNX 란?: https://wooono.tistory.com/415
• [Blog]PyTorch 2.0 vs ONNX vs TensorRT 비교: https://thecho7.tistory.com/entry/PyTorch-20-vs-ONNX-vs-TensorRT-%EB%B9%84%EA%B5%90
• [Blog] ONNX(Open Neural Network Exchange) 이해하기 -1: React Native 활용: https://adjh54.tistory.com/203