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 제품이 아닐 수도 있습니다. 잘 고려하셔야 합니다.