PyTorch 2.0 vs ONNX vs TensorRT 비교

안녕하세요, 딥러닝 모델을 개발하면 그 모델을 배포하는 과정이 매우 중요합니다. 어떻게 하면 더 빠르게 동작하는 서비스를 만들까 고민이 많으실텐데요, 오늘은 그 방법 중 몇가지를 비교하고 더 좋은 방법을 찾아보고자 글을 쓰도록 하겠습니다. 방법으로는 꽤 많이 알려진 PyTorch, ONNX, TensorRT를 골랐습니다. 왜 JAX가 없냐 물어보시면 JAX는 독자적 생태계가 확고해서 다음에 따로 리뷰하도록 하겠습니다. 참고로 글의 내용은 Nebuly Blog를 참조했습니다.

---

비교에 앞서 PyTorch, ONNX, TensorRT가 무엇인지 짧게 소개하면 이해가 더 쉬울 것 같습니다.

PyTorch 2.0

언젠가부터 Tensorflow를 앞지르고 가장 인기있는 Deep Learning Framework이 된 PyTorch가 2022년 12월에 새로운 버전을 공개했습니다.

이 글에서 알아야 할 중요한 부분은 PyTorch는 너무 Pythonic해서 문제라는 것입니다.

Python의 특징

그럼 Pythonic 하다는 게 뭐냐?

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

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

Pytorch 2.0의 특징

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

ONNX

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

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

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

 

요약하면 우리 모두가 C/C++/JAVA 가 파이썬보다 빠른 것은 알고 있죠? 그럼 파이썬에서 만든 모델을 C로 실행할 수 있다면 너무 좋겠네요. 그래서 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/

---

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 (Mixed Precision 또는 FP16이 좋아보임)
2. 추론할 때는 TensorRT가 좋아보이긴 함. 하지만 딥러닝 모델을 배포할 때는 GPU를 쓰지 않을 수도, GPU가 NVIDIA 제품이 아닐 수도 있습니다. 잘 고려하셔야 합니다.