로봇 ML 모델의 경량화 1부: 훈련 후 양자화

Aug.22.2024 박준수

Deep Learning Machine Learning Robotics

오늘날 머신러닝(machine learning, ML) 모델 개발은 대부분 NVIDIA의 데이터센터 GPU(예: A100)나 워크스테이션 GPU(예: RTX4090)가 장착된 고성능 서버 환경에서 이루어집니다. 우아한형제들 로보틱스LAB에서도, 실외 배달 로봇의 자율주행에 사용할 ML 모델을 개발할 때 이러한 고성능 서버들을 사용합니다. 덕분에 매우 큰 데이터셋들과 다양한 고성능 ML 모델들을 손쉽게 다루고 있죠. 그러나 여기엔 한 가지 문제점이 있는데, 바로 고성능 서버 환경에서 개발된 ML 모델은 곧바로 로봇에 배포할 수 없다는 점입니다.

이 글에서는 고성능 서버 환경과 실외 자율주행 로봇 환경의 차이점을 살펴보고, 그로부터 ML 모델 경량화의 필요성을 이해한 후, ML 모델 경량화 방법 중 하나인 양자화의 원리와 적용 방법을 알아보겠습니다. 양자화의 원리와 적용 방법에 대해서는 NVIDIA의 공식 문서를 참고하였으며, 이러한 기술들은 로봇을 포함한 다양한 NVIDIA 하드웨어 기반 컴퓨터에 응용할 수 있습니다.

로봇이 실외에서 자율주행을 하려면?

Throughput

그림 1. 우아한형제들 로보틱스LAB의 실외 자율주행 배달로봇 ‘딜리’

로봇이 실외 환경에서 자율주행을 하기 위해서는 무엇을 고려해야 할까요? 별도의 서버실이나 데이터센터에서 관리되는 고성능 서버와 비교해 보면, 실외 자율주행 로봇에는 다음과 같은 특성들이 추가로 요구됩니다.
  • 충격, 진동, 온도, 습도, 물, 먼지 등을 견디는 내구성
  • 긴 배터리 수명과 낮은 발열을 위한 전성비
  • 좁은 로봇 내부 공간에 맞는 작은 크기의 부품들
  • 다양한 센서를 연결하기 위한 많은 포트높은 호환성
이러한 조건을 무시하고 고성능 서버 환경의 GPU를 로봇에 그대로 사용하면, 조금의 충격에도 쉽게 망가지고 배터리도 금방 방전되는 로봇이 될 것입니다. 따라서 실외 환경에서 사용할 로봇에는 이러한 요구 사항들을 충족하는 엣지 디바이스(edge device)를 사용해야 합니다.
일반적으로 ML 목적의 엣지 디바이스에는 행렬 연산에 최적화된 처리 장치를 사용합니다. 그러한 처리 장치로 GPU를 사용할 수도 있고, GPU 대신 포괄적 개념의 ML 모델 처리 장치인 신경망 처리 장치(neural processing unit, NPU), Google의 텐서 처리 장치(tensor processing unit, TPU), Intel의 비전 처리 장치(vision processing unit, VPU), NVIDIA의 딥러닝 가속기(deep learning accelerator, DLA) 등을 사용하기도 합니다. 이러한 처리 장치들은 엣지 디바이스에서 GPU의 역할을 대체합니다.
그렇다면 수많은 ML 목적의 엣지 디바이스 중에서 어떤 제품을 선택해야 할까요? 우아한형제들 로보틱스LAB에서는 NVIDIA의 Jetson 플랫폼을 사용하고 있습니다. 다음 챕터에서는 NVIDIA GPU와 Jetson 플랫폼의 특징을 살펴보고, 그로부터 Jetson 플랫폼의 장점을 알아보겠습니다.

NVIDIA GPU와 Jetson 플랫폼의 특징

NVIDIA GPU는 ML 모델을 처리하는 장치로 많이 쓰입니다. 어떤 ML 처리 장치를 사용하든, 그것으로 ML 모델을 구동시키려면 딥러닝 프레임워크(PyTorch 등)의 포맷으로 되어 있는 ML 모델을 그 처리 장치의 특성에 맞게 최적화된 포맷으로 변환하는 과정이 필요합니다. 만약 ML 모델 내에 그 장치가 지원하지 않는 연산자(operator)가 있으면, 직접 커스텀 연산자를 만들어야만 변환을 할 수 있습니다.
NVIDIA GPU는 TensorRT를 통해 모델의 최적화 및 변환을 수행할 수 있습니다. TensorRT는 NVIDIA GPU를 위한 ML 모델 및 추론을 최적화할 때 쓰이는 도구이자 모델 추론을 수행하는 런타임 라이브러리입니다. ML 모델이 TensorRT의 최적화에 의해 변환된 결과를 TensorRT 엔진이라고 부릅니다. NVIDIA GPU의 코어인 CUDA 코어가 사용된 모든 하드웨어는, TensorRT 런타임 라이브러리를 사용하여 TensorRT 엔진을 추론할 때 가장 좋은 성능을 보입니다.

NVIDIA에서 공개한 문서 “NVIDIA TensorRT – Inference 최적화 및 가속화를 위한 NVIDIA의 Toolkit”에 따르면, TensorRT는 양자화(quantization), 그래프 최적화(graph optimization), 커널 자동 튜닝(kernel auto-tuning), 동적 텐서 메모리(dynamic tensor memory), 다중 스트림 실행(multi-stream execution)과 같은 기법들을 통해 모델의 추론 속도를 향상시킵니다. 이 문서에 제시된 실험 결과에 따르면 이러한 TensorRT의 최적화만으로도 ResNet-50 모델의 성능이 8배 이상 향상되었다고 합니다. NVIDIA는 다른 문서 “TensorRT Integration Speeds Up TensorFlow Inference”에서도 TensorRT를 사용했을 때의 성능 향상이 [그림 2]의 그래프처럼 매우 큼을 보였습니다. TensorRT가 이처럼 큰 차이를 만들기 때문에, 배포 단계에서 모델을 TensorRT 엔진으로 최적화하는 일은 필수입니다.

Throughput

그림 2. TensorFlow vs TensorRT ResNet-50 성능 비교 (출처: NVIDIA)

Jetson 플랫폼에서는 ML 모델 처리 장치로 CUDA 코어와 DLA를 모두 사용할 수 있습니다. DLA는 저전력으로 ML 모델의 연산을 수행하는 것에 특화된 처리 장치이며, CUDA 코어와 동일한 소프트웨어를 사용합니다. 이러한 특징 때문에, 고성능 서버 GPU에서 쓰이는 TensorRT와 같은 ML 모델 최적화 도구 및 방법들을 Jetson 플랫폼에서도 그대로 쓸 수 있습니다. 이로 인해 Jetson 플랫폼은 다음과 같은 이점을 가집니다.
1. NVIDIA의 고성능 서버 GPU에서 개발된 오픈 소스 모델들을 Jetson 플랫폼에서도 손쉽게 사용할 수 있습니다. 오늘날의 오픈 소스 모델들은 TensorRT 엔진으로의 변환을 잘 지원하고, 이에 필요한 커스텀 연산자를 지원하는 경우도 많습니다. 덕분에, Jetson 플랫폼에서는 오픈 소스 모델들을 이용하여 모델을 빠르게 개발하고 테스트할 수 있습니다.
2. CUDA와 TensorRT에 대한 문서화와 커뮤니티 활성화가 잘 되어 있습니다. 이로 인해 새로운 기능 개발, 모델 최적화, 이슈 해결 등에 필요한 정보를 쉽게 탐색할 수 있습니다.
3. CUDA의 시장 점유율이 높아서, Jetson 플랫폼 기반의 엣지 디바이스 제품군이 많이 있습니다. 따라서 다양한 공급망을 쉽게 구축할 수 있어 제품 단가를 절감할 수 있고 유지 보수를 하기에도 좋습니다.
이러한 장점들 때문에, 우아한형제들 로보틱스LAB을 포함한 많은 기업과 연구소들은 고성능 서버 환경과 엣지 디바이스 환경 모두에서 NVIDIA의 제품을 사용하고 있습니다.
Jetson 플랫폼 덕분에 로봇이 실외 환경에서도 안정적으로 오랜 시간 동작할 수 있게 되었지만, 그 대가로 연산 능력이 희생되었습니다. 따라서, 자율주행에서 요구되는 실시간성을 충족하기 위해서는 로봇에 배포될 ML 모델을 경량화하는 것이 필수적입니다.

ML 모델의 경량화는 크게 모델링(modeling) 및 훈련(training) 단계에서의 경량화, 그리고 훈련 완료 후의 경량화로 나눌 수 있습니다. 다음 챕터에서는 훈련 완료 후의 경량화 방법인 양자화에 대해 알아보겠습니다.

양자화

양자화(quantization)는 ML 모델의 연산 과정에서 사용되는 가중치(weight)와 출력 텐서(tensor)의 데이터 타입을 높은 정밀도(precision)에서 낮은 정밀도로 변환하는 경량화 방법입니다. 일반적으로 32비트(FP32) 또는 16비트(FP16) 값을 8비트(INT8, FP8) 또는 4비트(INT4) 값으로 변환하며, 주로 정숫값으로의 변환(INT8, INT4 등)을 우선적으로 고려합니다. 이렇게 정숫값으로 양자화된 텐서는 행렬 곱셈(matrix multiplication) 등의 행렬 연산을 정수 연산만으로 수행할 수 있게 됩니다. 정수 행렬 연산은 실수 행렬 연산에 비해 연산 속도가 매우 빠르기 때문에 전체 추론 속도가 크게 향상됩니다. 또한, 데이터가 16비트 이상의 정밀도에서 8비트 이하의 정밀도로 바뀌면서 모델의 크기도 줄어드는 효과를 얻을 수 있습니다. 하지만 표현할 수 있는 값의 범위도 줄어들기 때문에, 양자화 전후 값의 매핑 효율이 낮을수록 정확도 손실도 커집니다. 매핑 효율을 높이기 위해서는 캘리브레이션(calibration)을 통해 타깃 데이터와 모델의 특성에 맞는 최적의 양자화 계수를 구해야 합니다.

NVIDIA의 양자화 관련 문서 “Achieving FP32 Accuracy for INT8 Inference Using Quantization Aware Training with NVIDIA TensorRT”에 따르면, TensorRT에서 FP32를 INT8로 양자화할 때 캘리브레이션을 수행하는 과정은 다음과 같습니다.

1. 먼저 FP32 값과 INT8 값을 매핑하기 위한 양자화 수식을 설정합니다. 문서에서는 아래 수식과 같이 선형 관계로 표현하고 있습니다. 여기서 Q는 양자화된 텐서, X는 원본 텐서, s는 스케일(scale)입니다.

Quant formula

2. 양자화하고자 하는 텐서의 분포를 구합니다. 텐서의 분포는 일반적으로 모델 내부의 다양한 정규화(normalization) 연산에 의해 [그림 3]과 같이 0을 중심으로 하는 정규 분포가 됩니다.

3. 분포 내의 모든 실숫값을 정숫값에 1대 1 매핑할 수 없으므로, 매핑 범위의 임곗값([그림 3]에서 amax)을 설정합니다. 그리고 절댓값이 해당 임곗값을 넘는 모든 값을 임곗값으로 바꿔줍니다(이 연산을 Clip이라고 합니다). 이를 양자화 수식에 반영하면 아래와 같습니다.

Quant formula

4. 임곗값 이내의 범위인 [-amax, amax] 내의 실수와 정수 데이터 타입의 전체 범위(예: INT8이라면 [-128, 127])를 매핑할 수 있는 스케일값을 구합니다. 이를 수식으로 나타내면 아래와 같습니다.

Quant formula

5. 충분한 양의 학습 데이터에 대해 위의 과정을 반복합니다. 이를 통해, 타깃 데이터를 양자화하기 전과 후 사이의 오차를 가장 많이 줄여주는 임곗값과 스케일값을 찾고, 이 값들을 저장합니다.

Quant formula

그림 3. 캘리브레이션의 원리 (출처: NVIDIA)

학습이 완료된 모델을 대상으로 캘리브레이션 및 양자화를 수행하는 과정을 훈련 후 양자화(post-training quantization, PTQ)라고 합니다. 훈련 후 양자화는 추가 훈련 과정이 필요하지 않고 소량의 학습 데이터만 있으면 되기 때문에, 쉽게 수행해 볼 수 있습니다. 또한 다른 경량화 방법들과 비교했을 때, 잘못된 튜닝으로 인해 정확도가 크게 하락할 가능성이 적습니다.
만약 훈련 후 양자화를 적용한 결과 정확도 손실이 크다면, 양자화 인식 훈련(quantization aware training, QAT)을 고려해 볼 수 있습니다. 양자화 인식 훈련은 양자화의 영향을 모델의 가중치 튜닝에 반영하는 방법입니다. 이 방법으로 캘리브레이션 과정에서 발생하는 정확도 손실을 보정할 수 있지만, 여기에는 전체 학습 데이터를 사용하는 추가 훈련 과정이 필요합니다.

이 글에서는 훈련 후 양자화 방법부터 자세히 다루고, 양자화 인식 훈련에 대한 자세한 내용은 “로봇 ML 모델의 경량화” 시리즈의 다음 글에서 다룰 예정입니다.

TensorRT를 이용한 최적화 방법

이번 챕터에서는, 고성능 서버 환경에서 딥러닝 프레임워크를 이용해 학습한 모델을 TensorRT 엔진으로 변환하는 과정을 예시 코드와 함께 살펴보겠습니다.
예시에서는 PyTorchResNet-18 모델과 Hugging Face의 ImageNet validation 데이터셋을 사용하였으며, [표 1]과 같은 환경에서 테스트하였습니다.

Quant formula

표 1. 예시 코드의 테스트 환경

PyTorch 모델을 TensorRT로 추론하는 방법 비교

현재 TensorRT는, PyTorch 모델을 직접 TensorRT 엔진으로 변환하는 도구를 지원하지 않습니다. 대신, PyTorch 모델을 TensorRT 런타임으로 추론할 수 있는 두 가지 방법이 있습니다. 첫 번째는 ONNX 모델로 변환 후 다시 TensorRT 엔진으로 변환하는 방법이고, 두 번째는 TorchScript로 변환 후 Torch-TensorRT를 이용하는 방법입니다. 이 두 방법의 특징은 아래와 같습니다.

1. ONNX 모델 변환 후 TensorRT 엔진 변환

이 방법은 PyTorch로 훈련된 모델을 ONNX 모델로 변환한 뒤, 이를 다시 TensorRT 엔진으로 변환하여 사용하는 것입니다. ONNX는 ML 모델 포맷의 한 종류인데, 대부분의 딥러닝 프레임워크가 ONNX로의 변환 또는 ONNX로부터의 변환을 지원합니다. PyTorch도 ONNX로의 변환을 지원하며, TensorRT는 ONNX 모델로부터 TensorRT 엔진을 빌드할 수 있습니다. 이 방법의 장점은 최적화가 잘 된 온전한 TensorRT 엔진을 사용할 수 있다는 것입니다. 그러나, PyTorch 모델에 ONNX나 TensorRT가 지원하지 않는 연산자가 포함된 경우에는 직접 커스텀 연산자를 만들어야만 변환이 가능하다는 단점이 있습니다.

2. Torch-TensorRT를 이용한 추론

Torch-TensorRT는 TorchScript 모델을 TensorRT를 이용해 최적화할 수 있도록 해주는 컴파일러입니다. TorchScript는 PyTorch 모델의 파이썬(Python)에 대한 의존성(dependency)을 없애고, 다양한 환경에서 고성능으로 추론할 수 있도록 모델을 컴파일한 포맷입니다. TorchScript 모델을 Torch-TensorRT를 이용해 변환하면 TensorRT와 호환되는 연산자는 TensorRT 연산자로 변환됩니다. 이렇게 변환된 연산자들은 TensorRT의 최적화 기능들을 활용할 수 있습니다. 한편, TensorRT와 호환되지 않는 연산자는 PyTorch 연산자가 처리합니다. 따라서, 이 방법은 ONNX를 통한 변환과 달리 호환성 문제가 적습니다.
변환된 TorchScript 모델 내에서 PyTorch 연산자와 TensorRT 연산자가 함께 사용될 경우, [그림 4]와 같은 연산 과정을 거칩니다. 먼저, PyTorch 연산자의 출력이 TensorRT 연산자와 만나면 TorchScript 인터프리터는 TensorRT 엔진을 호출하여 모든 입력을 전달합니다. 호출된 TensorRT 엔진은 해당 입력에 대한 연산을 수행한 뒤, 연산 결과를 TorchScript 인터프리터에 다시 전달합니다.

Quant formula

그림 4. Torch-TensorRT의 런타임 구조 (출처: NVIDIA)

이처럼 Torch-TensorRT는 TorchScript를 기반으로 하기 때문에, 온전한 TensorRT 엔진에 비해 추론 속도가 느리며, 원래의 TorchScript 모델 내부의 연산자 중 TensorRT 호환 연산자가 적을수록 추론 속도는 더욱 느려집니다.

자율주행과 같이 추론 속도가 매우 중요한 분야에서는 온전한 TensorRT 엔진을 사용하는 “ONNX 모델 변환 후 TensorRT 엔진 변환” 방식을 추천드립니다. 이 글의 예시 코드에도 그러한 방식을 사용하였습니다.

ONNX 모델 변환

PyTorch는 모델을 ONNX 포맷으로 변환할 수 있는 내장 함수를 제공하고 있습니다.
아래는 Torchvision의 ResNet-18 ImageNet 사전훈련(pretrained) 모델을 선언하고 이를 ONNX 모델로 변환하는 예시입니다.
import torch
import torchvision

# Load ResNet-18 ImageNet-pretrained model using torchvision.
model = torchvision.models.resnet18(pretrained=True)
model.eval()

# Dummy input with the model input size.
dummy_input = torch.randn(1, 3, 224, 224)

# Save ONNX model to ONNX_PATH.
torch.onnx.export(model, 
                  dummy_input, 
                  ONNX_PATH, 
                  opset_version=13,
                  input_names=["input"], 
                  output_names=["output"])
위와 같이 torch.onnx.export를 이용하면 PyTorch의 모델을 간단하게 ONNX 모델로 변환할 수 있습니다. 각 인자의 역할은 다음과 같습니다.
  • model: 변환할 PyTorch 모델 객체를 입력합니다.
  • dummy_input: PyTorch 모델의 입력과 같은 크기의 임의의 텐서를 입력합니다.
  • ONNX_PATH: ONNX 모델을 저장할 경로를 지정합니다.
  • opset_version: ONNX의 연산자 버전을 지정합니다. 자신이 사용할 TensorRT 버전이 지원하는 범위 내의 버전으로 입력해 주어야 합니다.
  • input_names: ONNX 포맷에서 사용되는 입력의 이름을 지정합니다. 입력이 여러 개일 경우를 위해 list 형태로 입력합니다. 지정된 이름은 추론 단계에서 입력 dictkey로 사용됩니다.
  • output_names: ONNX 포맷에서 사용되는 출력의 이름을 지정합니다. 출력이 여러 개일 경우를 위해 list 형태로 입력합니다. 지정된 이름은 추론 단계에서 출력 dictkey로 사용됩니다.

만약 PyTorch 모델에 ONNX에서 지원하지 않는 레이어가 포함되어 있다면 이러한 변환이 불가능하므로, ONNX 공식 문서를 참고하여 커스텀 레이어를 만들어 사용해야 합니다.

Polygraphy를 이용한 훈련 후 양자화 및 TensorRT 엔진 변환

Polygraphy는 NVIDIA에서 제공하는 개발 도구입니다. 이것은 TensorRT API를 기반으로 만들어졌고, 고급(high-level) 추상화와 인터페이스를 제공합니다. 따라서 TensorRT API보다 사용하기 간편하고, 버전 변경에 따른 코드 호환성 문제도 훨씬 적습니다. 또한, Polygraphy로 생성된 객체가 TensorRT API에 기반하므로 Polygraphy와 TensorRT API를 혼용하여 사용할 수 있습니다.

이번 챕터에서는 Polygraphy를 이용한 훈련 후 양자화 및 TensorRT 엔진으로 변환하는 방법을 알아보겠습니다.

1. Calibrator 객체 생성

훈련 후 양자화를 수행하려면 Polygraphy의 Calibrator 클래스를 사용해야 합니다. 이 Calibrator 클래스는 입력 데이터를 generator의 형태로 받습니다. 따라서 추론 환경과 동일한 전처리가 적용된 입력 데이터를 생성하는 generator를 만들고, 이를 Calibrator에게 인자로 전달해야 합니다.
아래는 generator 정의 및 Calibrator 객체 생성 예시입니다.
import os

from PIL import Image
from polygraphy.backend import trt as poly_trt
from torchvision import transforms

# Preprocess of ResNet-18 training.
transform = transforms.Compose([transforms.Resize((256, 256)), 
                                transforms.CenterCrop((224, 224)),
                                transforms.ToTensor(), 
                                transforms.Normalize([0.485, 0.456, 0.406], 
                                                     [0.229, 0.224, 0.225])])

# Polygraphy needs a generator-type data loader.
val_list = os.listdir(IMAGE_DIR)
def data_generator():
    for image in val_list:
        # Preprocess.
        image_path = os.path.join(IMAGE_DIR, image)
        image = transform(Image.open(image_path).convert("RGB"))
        # Add batch dimension.
        image = image.unsqueeze(0)
        # Polygraphy uses numpy input.
        image = image.numpy()
        # Dict key must be the same as ONNX input name.
        yield {"input": image}

calibrator = poly_trt.Calibrator(data_loader=data_generator())
data_generatorIMAGE_DIR에서 입력 이미지를 읽어와 ResNet-18과 동일한 전처리 과정을 수행한 후, 앞서 만든 ONNX 모델의 입력 형식과 동일한 형식의 dict를 반환합니다. 그리고 Calibrator 타입의 객체를 만드는데, 이때 위에서 만든 generator를 인자로 사용합니다.

2. ONNX 파일 로드 및 IBuilderConfig 객체 생성

Calibrator를 생성한 후에는, ONNX 모델을 로드하고, TensorRT 엔진 빌드에 필요한 각종 설정을 담고 있는 IBuilderConfig 객체를 만들어야 합니다.
아래는 ONNX 모델 로드 및 IBuilderConfig 객체 생성 예시입니다.
builder, network, parser = poly_trt.network_from_onnx_path(path=ONNX_PATH)

# Each type flag must be set to true.
builder_config = poly_trt.create_config(builder=builder,
                                        network=network,
                                        int8=True,
                                        fp16=True,
                                        calibrator=calibrator)
이처럼 network_from_onnx_path를 이용해 ONNX 모델을 로드하면 builder, network, parser 세 가지 객체가 반환됩니다. 이 중 buildernetworkIBuilderConfig 객체를 만드는 create_config의 인자로 사용됩니다.
IBuilderConfig 객체를 만들 때는 정수 타입을 지원하지 않는 레이어를 위해 FP16 타입 변환에 대한 옵션을 추가해야 합니다. 여기서 주의할 점은, IBuilderConfig은 각 타입에 대한 인자를 해당 타입 사용 여부를 나타내는 플래그로 사용한다는 것입니다. 따라서, INT8과 FP16 타입을 모두 사용하려면 create_config에서 이들 타입 각각에 해당되는 인자를 모두 True로 설정해야 합니다.

마지막으로, 앞서 만든 Calibrator 객체를 인자로 추가하면 IBuilderConfig 객체 생성이 완료됩니다.

3. TensorRT 엔진 빌드 및 저장

아래는 TensorRT 엔진 빌드 및 저장 방법에 대한 예시입니다.
engine = poly_trt.engine_from_network(network=(builder, network, parser),
                                      config=builder_config)

# TensorRT engine will be saved to ENGINE_PATH.
poly_trt.save_engine(engine, ENGINE_PATH)
앞서 ONNX 모델로부터 로드된 builder, network, parser 객체들과 IBuilderConfig 객체를 engine_from_network에 인자로 넣으면, TensorRT 엔진이 빌드됩니다.

마지막으로, 빌드된 TensorRT 엔진 객체를 save_engine을 이용하여 저장하면 TensorRT 엔진 변환 과정이 완료됩니다.

4. 추론 방법

TensorRT 엔진을 저장한 후엔, 아래의 예시와 같이 저장된 엔진을 로드하여 추론 테스트를 할 수 있습니다.
# Load serialized engine using 'open'.
engine = poly_trt.engine_from_bytes(open(ENGINE_PATH, "rb").read())

with poly_trt.TrtRunner(engine) as runner:
    # Preprocess.
    image = transform(Image.open(IMAGE_PATH).convert("RGB"))
    image = image.unsqueeze(0).numpy()
    # Input dict keys are the same as 'input_names' arg in 'torch.onnx.export'.
    output_dict = runner.infer({"input": image})
    # Output dict keys are the same as 'output_names' arg in 'torch.onnx.export'.
    output = output_dict["output"]
먼저 engine_from_bytes를 이용해 TensorRT 엔진을 로드한 후, TrtRunner를 사용하여 로드된 엔진을 추론할 수 있는 runner 객체를 생성합니다. 이 runner 객체는 추론 시 runner.infer를 사용하며, 입력과 출력 모두 dict 타입을 사용합니다. 이때 입력 dictkey로는 ONNX 모델 생성 시 torch.onnx.export에 인자 input_names로 전달했던 값을, 출력 dictkey로는 인자 output_names로 전달했던 값을 사용해야 합니다. 즉, 이미지를 전처리한 후 {input_names: image} 형태의 dict로 변형하여 runner.infer에 입력합니다. 또한, 출력 dictkeyoutput_names를 이용해 value에 접근하여 모델의 출력을 얻을 수 있습니다.

양자화 단계별 성능 비교

이번 챕터에서는 양자화가 모델의 정확도, 추론 속도, 모델 크기에 미치는 영향을 확인하기 위해 어떤 방법으로 실험하였는지, 그 결과는 어땠는지를 살펴보겠습니다. 이 실험에서도 예시 코드와 마찬가지로 ResNet-18 모델과 ImageNet validation 데이터셋을 사용하였습니다.

비교군 설정

양자화된 모델에 2개의 비교군 모델을 추가하여 총 3개의 모델을 실험에 사용하였습니다. 비교군 모델들을 변환한 방식은 다음과 같습니다.

1. FP32 원본 모델

이것은 원본인 PyTorch ResNet-18 모델을 FP32 타입의 TensorRT 엔진으로 단순 변환한 모델입니다.
이를 위해 예시 코드에서 변경한 부분은 아래와 같습니다.
builder_config = poly_trt.create_config(builder=builder,
                                        network=network)
이처럼 create_config에서 정밀도 및 Calibrator 관련 인자를 제거하여 FP32 타입만으로 모델을 빌드하도록 변경하였습니다.

2. 랜덤 데이터를 Calibrator에 통과시킨 INT8 모델

캘리브레이션이 미치는 영향을 보기 위해, 원래의 학습 데이터 대신 랜덤 데이터 입력을 사용하고 INT8로 변환한 모델입니다.
이를 위해 예시 코드에서 변경한 부분은 아래와 같습니다.
# Generate random data instead of real data.
def data_generator():
    for _ in range(1):
        image = torch.randn(1, 3, 224, 224).numpy()
        yield {"input": image}

# Set calibrator with a new data generator.
calibrator = poly_trt.Calibrator(data_loader=data_generator())
이처럼 data_generator가 원래의 학습 데이터 대신 torch.randn으로 생성한 랜덤 데이터를 반환하도록 하였습니다.

실험 방식

테스트를 하기 위해 Polygraphy의 추론 기능을 활용하였으며, ImageNet validation 데이터셋의 모든 이미지에 대한 평균 추론 시간과 정확도를 측정하였습니다. 데이터의 구조 및 정답(label) 체크 방식에 대해서는 Hugging Face의 ImageNet 데이터셋을 참고해 주세요.
실험 코드는 아래와 같습니다.
import os
import time

from PIL import Image
from polygraphy.backend import trt as poly_trt
import tqdm

# Load TensorRT engine from ENGINE_PATH.
engine = poly_trt.engine_from_bytes(open(ENGINE_PATH, "rb").read())

# Refer to Hugging Face for IMAGENET2012_CLASSES.
class_list = list(IMAGENET2012_CLASSES.keys())
val_list = os.listdir(IMAGE_DIR)

total_infer_time = 0.0
total_correct = 0
with poly_trt.TrtRunner(engine) as runner:
    for image in tqdm.tqdm(val_list):
        # Get label from file name.
        label = image.split("_")[-1].split(".")[0]
        # Preprocess.
        image_path = os.path.join(IMAGE_DIR, image)
        image = transform(Image.open(image_path).convert("RGB"))
        image = image.unsqueeze(0).numpy()
        # Get inference time using 'time' module.
        before_infer = time.time()
        output_dict = runner.infer({"input": image})
        total_infer_time += time.time() - before_infer
        output = output_dict["output"]
        predicted = output.argmax()
        # Calculate total number of correct outputs.
        total_correct += class_list[predicted.item()] == label

# Calculate and print average values.
accuracy = 100 * total_correct / len(val_list)
avg_infer_time_ms = 1000 * total_infer_time / len(val_list)
print(f"Accuracy: {accuracy:.2f} %")
print(f"Average inference time: {avg_infer_time_ms:.2f} ms")
이 코드에서는 데이터가 들어있는 디렉터리 내의 모든 이미지 파일 각각에 대해 추론에 소요된 시간과 정답 여부를 계산하여 각각 total_infer_timetotal_correct에 누적합니다. 추론 시간은 runner.infer 실행 전후에 time.time을 사용하여 계산하였고, 정답 여부는 class_list에서 인덱스로 구한 클래스 이름이 이미지 이름에 포함된 클래스 이름과 일치하는지 체크하여 알아내도록 하였습니다.

모든 합산이 끝난 후엔 누적된 값들을 이미지 개수로 나누어 평균 추론 시간과 정확도를 출력하도록 하였습니다.

실험 결과 및 결론

실험 결과는 아래의 [표 2]와 같습니다.

Quant formula

표 2. 양자화 단계별 성능 비교 실험 결과

실험 결과에 따르면, INT8 모델들은 FP32 원본 모델에 비해 추론 시간에서 약 3배, 모델 크기에서 약 4배의 성능 향상을 보였습니다. 또한 랜덤 데이터 캘리브레이션이 적용된 INT8 모델은 FP32 원본 모델에 비해 3.25%의 정확도 손실이 있는 반면, 학습 데이터 캘리브레이션이 적용된 INT8 모델은 정확도 손실이 0.22%에 불과했습니다.

이로부터, Polygraphy를 이용한 학습 데이터 기반 캘리브레이션 및 양자화는 모델의 정확도를 거의 유지하면서도 추론 속도를 크게 향상시키고 모델 크기를 줄일 수 있음을 알 수 있습니다.

마치며

로봇의 연산 장치는 여러 제약으로 인해 성능이 낮기 때문에, 이를 해결하기 위해 다양한 ML 모델 최적화 기술이 필요합니다. 그러나 훈련 단계와 추론 단계에서 사용하는 하드웨어와 ML 모델의 포맷이 다르기 때문에, 배포 과정에서 고성능 서버 환경에 비해 더 많은 시행착오를 겪게 됩니다.
오늘 소개한 방법을 사용하면, 이러한 시행착오를 최소화하고 어떤 모델이든 빠르게 경량화 및 TensorRT 엔진 배포를 할 수 있습니다. 이 방법은 로봇뿐만이 아닌 TensorRT를 이용한 배포가 필요한 다양한 환경에서 모두 사용할 수 있습니다. 비슷한 문제를 겪는 분들께 이 글이 도움이 되었기를 바랍니다.

로봇 ML 모델 경량화 시리즈의 두 번째 글에서는 모델 훈련 과정에서 사용할 수 있는 또 다른 양자화 방법인 양자화 인식 훈련에 대해 다룰 예정이니 많은 관심 부탁드립니다.