로봇 ML 모델의 경량화 2부: 양자화 인식 훈련

1 month ago 8

"로봇 ML 모델의 경량화 1부: 훈련 후 양자화" 편에서는 Jetson 플랫폼의 특징, 양자화와 훈련 후 양자화에 대한 개념, TensorRT를 이용해 최적화를 하는 방법과 그 결과까지 알아봤습니다.

이번 편에서는 양자화로 인한 성능 저하를 막는 또 하나의 방법, 양자화 인식 훈련(Quantization Aware Training, QAT)에 대해 알아보겠습니다. 본격적인 내용에 앞서, 우아콘에서 발표한 "자율주행 로봇을 위한 머신러닝 모델의 추론 성능을 최적화하기" 영상을 시청하시면 이 글의 내용을 이해하는 데 도움이 될 수 있습니다.

훈련 후 양자화(Post-Training Quantization, PTQ)는 추가 훈련 없이 모델을 양자화해 경량화와 추론 속도를 개선하는 간단하고 효율적인 방법입니다. 하지만 PTQ에는 몇 가지 한계점이 있습니다. 아래에서 이에 대해 살펴보겠습니다.

오차 전파(에러 프로파게이션)에 대한 대응 부족

PTQ 모델은 추론 과정에서 각 레이어의 양자화 오차가 다음 레이어로 전파되어 누적될 수 있습니다. 실제 부동소수점 연산과 양자화 연산 사이의 오차는 예측 과정을 거치면서 점점 더 커질 수 있으며, 이는 모델 성능 저하로 이어집니다. 특히 깊은 신경망 구조에서는 이러한 오차 누적 현상이 더욱 심각해질 수 있습니다.

모델 구조 및 연산 특성 반영 부족

PTQ는 이미 학습이 완료된 모델 파라미터를 고정한 상태에서 사후 처리를 통해 양자화 범위(스케일, 제로 포인트 등)를 결정하게 됩니다. 여기서 모델 내부의 비선형 연산(예: 활성 함수, 배치 정규화 등)이나 레이어 간 분포 변화 등 복잡한 구조적 특성을 충분히 반영하기가 어렵습니다. 결국, 특정 레이어 또는 채널별 분포 차이를 세밀하게 고려하지 못하여 정밀도 손실이 심해질 가능성이 있습니다.

예시: 3D 객체 검출 모델에서 성능 저하

3D 객체 검출은 자율주행에서 중요한 역할을 수행하는 기술입니다. 2D 카메라 이미지에서 객체의 위치와 크기를 나타내는 3D 경계 상자를 예측하여 주변 환경을 3차원적으로 인식할 수 있도록 합니다. 아래 그림은 3D 객체 검출 모델의 예시를 보여줍니다.

3D 객체 검출 모델은 복잡한 구조로 인해 PTQ만으로는 양자화로 인한 성능 저하를 완전히 막기 어렵습니다. 다양한 크기의 객체를 정확하게 검출하기 위해 여러 층의 컨볼루션 연산과 복잡한 특징 추출 과정을 거치는데, 이 과정에서 양자화 오차가 누적되어 성능에 영향을 미칠 수 있기 때문입니다.

실제로 3D 객체 검출 모델에 PTQ를 적용했을 때 성능 저하가 얼마나 발생하는지 확인하기 위해 FP32 모델을 기준으로 INT8 모델과 INT8+PTQ 모델의 성능을 비교하는 실험을 진행했습니다. 아래 테이블에서 FP32 는 32비트 부동 소수점을 사용하는 모델, INT8 은 랜덤 데이터 캘리브레이션이 적용된 INT8 모델, INT8+PTQ은 학습 데이터 캘리브레이션이 적용된 INT8 모델을 의미합니다. 각 모델의 성능은 mAP(mean Average Precision)를 사용하여 평가했으며, mAP는 객체 검출 모델의 성능을 평가하는 지표로 값이 클수록 성능이 우수함을 나타냅니다. 아래 표는 실험 결과를 보여줍니다.

[표 1]을 보면 FP32 모델의 mAP는 0.2834로 가장 높은 성능을 보입니다. 반면, INT8 양자화를 적용한 모델은 mAP가 0.1943으로 크게 감소하여 성능 저하가 뚜렷하게 나타났습니다. PTQ를 적용한 INT8+PTQ 모델은 mAP가 0.2435로 INT8 모델보다는 성능이 향상되었지만, FP32 모델에 비해서는 여전히 낮은 성능을 보입니다. 이는 PTQ가 양자화로 인한 성능 저하를 완전히 해결하지 못한다는 것을 의미합니다. 즉, 모델의 크기를 줄이고 연산 속도를 높이는 데에는 효과적이지만, 정확도를 유지하는 데에는 한계가 있다는 것을 알 수 있습니다.

이렇듯 PTQ는 훌륭한 방법론이지만 그 한계가 있는 것을 확인할 수 있습니다. 이러한 PTQ의 한계를 극복하기 위해 양자화 인식 훈련(Quantization Aware Training, QAT)이라는 방법을 씁니다. 이에 대해 자세히 알아보도록 하겠습니다.

QAT는 모델을 처음부터 양자화 연산을 고려하여 학습시키는 방법입니다. 즉, 훈련 과정에서부터 모델이 양자화된 상태임을 가정하고 학습하여 양자화로 인한 정밀도 손실을 최소화하는 방법입니다.

QAT는 훈련 단계에서 부동소수점 연산을 그대로 사용하지 않고, 활성 함수나 가중치 텐서 등에 양자화 연산을 모방한 기법을 적용합니다. 모델은 훈련 과정에서 양자화로 인해 발생하는 오차에 적응하게 되어, 실제 양자화된 환경에서도 높은 성능을 유지할 수 있습니다. QAT가 수행되는 원리를 이해하기 위해서는 fake quantization, 혹은 Q/DQ라 불리는 개념과 narrow minima, wide minima의 개념을 이해해야 합니다.

Fake Quantization, Q/DQ

QAT의 핵심 메커니즘이 바로 fake quantization이라고도 부르는 Q/DQ 연산입니다. 아래 [그림 2]를 보면 Layer1에서 FP32 상태로 출력된 값이 Quantize → Dequantize 과정을 거치면서 가상의 INT8 연산을 경험하게 되는 것을 알 수 있습니다.

Quantize (FP32 → INT8)

Quantize (FP32 → INT8) 단계에서는 FP32 텐서를 INT8 범위로 변환하여 의도적으로 오차를 발생시킵니다. 이때 실제로 메모리에 INT8 형태로 저장하는 것이 아니라, 값의 범위만 INT8처럼 제한하는 시뮬레이션을 수행합니다. 즉, 컴퓨터는 여전히 FP32 형태로 값을 저장하고 있지만, INT8처럼 작동하도록 범위를 제한하는 것입니다.

Dequantize (INT8 → FP32)

Dequantize (INT8 → FP32) 단계에서는 Quantize 단계에서 제한된 값을 다시 FP32 범위로 복원합니다. 역전파 시 부동소수점 정보가 그대로 유지되므로 학습 효율을 높일 수 있습니다.

Q/DQ 연산을 통해 모델은 순전파 시에는 INT8의 낮은 정밀도 환경을 경험하고, 역전파 시에는 FP32의 정확한 미분 값을 사용하여 가중치를 갱신합니다. 이처럼 Q/DQ는 모델이 정밀도 저하로 인한 오차에 적응하고, 양자화된 환경에서도 최적의 성능을 낼 수 있도록 돕는 역할을 합니다.

다음으로 알아볼 것은 Q/DQ를 이용해 낮은 정밀도를 체험할 때, loss가 수렴하는 모습입니다. 이를 위해 먼저 narrow minima, wide minima에 대한 개념을 정리할 필요가 있습니다.

Narrow Minima vs. Wide Minima

딥러닝 모델을 학습시킬 때, 목표는 손실 함수를 최소화하는 지점인 local minima를 찾는 것입니다. 양자화 과정에서는 손실 함수 값 이외에 local minima 주변 지형의 ‘기울기’도 중요한 역할을 합니다. 모델의 안정성이 달라지기 때문이죠.

아래 그림은 local minima의 두 가지 유형인 narrow minima와 wide minima를 보여줍니다.

Narrow Minima

Minima 주변이 급격하게 변하는 지형입니다.이러한 지형에서는 파라미터 값이 조금만 변해도 손실값이 크게 증가합니다. 예를 들어, FP32에서 INT8로 양자화하는 과정에서 발생하는 작은 변화에도 모델의 성능이 크게 저하될 수 있습니다. 따라서 narrow minima는 정밀도 저하나 노이즈에 매우 민감합니다.

Wide Minima

Minima 주변이 완만하게 변하는 지형입니다. 이러한 지형에서는 가중치가 약간 변하더라도 손실값은 크게 변하지 않습니다. 따라서 wide minima는 양자화나 노이즈에 대한 견고성(Robustness)이 높습니다. 즉, 외부 요인에 의해 파라미터 값이 조금 변하더라도 모델의 성능이 크게 영향을 받지 않습니다.

이제 Q/DQ와 이를 연관지어 설명해보겠습니다. FP32 가중치를 INT8과 같이 더 낮은 해상도로 표현하는 양자화 과정에서는 필연적으로 오차가 발생합니다. Q/DQ, 혹은 fake quantization 방법은 훈련 단계에서부터 이러한 오차를 미리 경험하도록 하여, 마치 가중치가 INT8로 양자화된 것처럼 값의 범위를 제한합니다.

제한된 범위의 가중치로 학습을 진행할 때 모델은 INT8 정밀도 내에서 손실을 최소화하려고 합니다. INT8의 한계로 인해 FP32에서 달성할 수 있는 최적의 성능에는 미치지 못할 수 있지만, 주어진 제약 조건 내에서 최적의 성능을 찾기 위해 노력합니다. 이러한 이유로, 일반적으로 minima는 wide minima 근처에 위치하게 됩니다. 값의 범위가 제한된 경우에는 wide minima의 완만한 기울기가 수렴에 더 유리하기 때문입니다.

위 [그림 2]를 예로 들어 설명을 이어가보겠습니다. 만약 weight 값을 INT8 형태로 표현한다면 그림에서 분홍색 포인트로 표현된 것 처럼 1, 2, 3, 4, 5와 같이 다섯 개의 값만 가질 수 있습니다. FP32를 사용한다면 weight 값은 4.2처럼 소수점까지 자유롭게 표현할 수 있기 때문에 narrow minima의 가장 낮은 지점에 도달할 수 있습니다. 하지만 INT8의 경우에는 표현할 수 있는 값이 제한적이기 때문에 wide minima에서 가장 가까운 값인 2를 선택하게 됩니다.

정리

지금까지 살펴본 QAT는 Q/DQ를 활용하여 낮은 해상도 환경을 미리 경험하면서 학습하는 기법입니다. 즉, INT8과 같은 낮은 해상도에서도 최소한의 손실을 갖도록 모델을 최적화하는 것이죠. QAT는 narrow minima가 아닌 wide minima 근처로 수렴하도록 유도하여 양자화로 인한 손실 증가를 최소화합니다.

QAT는 학습 후 모델을 변환하는 PTQ의 단점을 보완합니다. PTQ는 INT8로 변환할 때 narrow minima의 영향을 크게 받아 손실이 증가할 수 있는데, QAT는 이러한 문제를 미리 해결하여 안정적인 성능을 확보합니다.

하지만 QAT는 기존에 학습된 모델을 사용할 수 없고 처음부터 다시 학습해야 한다는 단점이 있습니다. 이는 추가적인 시간과 자원을 필요로 하므로, 특히 대규모 데이터셋이나 복잡한 모델에 적용할 때는 부담이 될 수 있습니다. 또한 모델이나 데이터의 규모가 작을 때에는 큰 효과를 발휘하지 못하기도 합니다.

다음 장에서는 QAT를 실제로 어떻게 수행하는지 단계별로 알아보겠습니다.

QAT를 수행하는 방법에는 크게 두 가지가 있습니다. 직접 Q/DQ 노드를 구현하는 방법과, 이미 구현된 라이브러리를 활용하는 방법이죠. 이 글에서는 PyTorch 기반, QAT를 지원하는 대표적인 두 가지 라이브러리 두 가지를 소개하고, 각각의 장단점을 비교해 보겠습니다.

PyTorch의 torch.ao.quantization

장점: PyTorch 생태계와 긴밀하게 통합되어 있으며, 다양한 모델에 적용할 수 있도록 커스터마이징 옵션을 풍부하게 제공합니다.
단점: 엣지 디바이스 배포에 많이 쓰이는 NVIDIA TensorRT와의 호환성이 떨어져 추가적인 변환 과정이 필요할 수 있습니다.

NVIDIA의 pytorch-quantization

장점: NVIDIA TensorRT와의 호환성이 뛰어나 엣지 디바이스를 포함한 NVIDIA 하드웨어 환경에 적합합니다. 또한 사용법이 간편하고 NVIDIA 생태계와 잘 통합되어 있습니다.
단점: torch.ao.quantization에 비해 커스터마이징 옵션이 제한적입니다.

이 글에서는 로봇 배포를 위해 NVIDIA의 pytorch-quantization 모듈을 이용하여 QAT를 수행하는 방법을 자세히 알아보겠습니다. pytorch-quantization은 TensorRT와의 호환성이 뛰어나고 사용이 간편하기 때문에 로봇 개발에 유용하게 활용될 수 있습니다.

NVIDIA pytorch-quantization을 활용한 QAT 수행 및 ONNX/TensorRT 변환 가이드

여기에서는 NIVDIA pytorch-quantization 모듈을 활용하여 QAT를 수행하고, 훈련된 모델을 ONNX 및 TensorRT로 변환하는 방법을 단계별로 설명합니다. ResNet18 모델과 CIFAR-10 데이터셋을 예시로 사용합니다.

1단계: 라이브러리 설치

먼저 다음 명령어를 실행하여 필요한 라이브러리를 설치합니다.

git clone https://github.com/NVIDIA/TensorRT.git cd TensorRT/tools/pytorch-quantization python3 setup.py install --user

2단계: pytorch-quantization 초기화

QAT를 수행할 수 있도록 pytorch-quantization 모듈을 초기화합니다.

from pytorch_quantization import quant_modules quant_modules.initialize()

pytorch-quantization 은 monkey patching 기법을 사용하여 Q/DQ 노드를 추가합니다. Monkey patching은 런타임에 코드를 수정하는 기법으로, 기존 레이어를 양자화된 연산을 포함하는 새로운 레이어로 교체합니다. quant_modules.initialize() 함수는 이러한 monkey patching을 수행합니다.

3단계: 데이터셋 준비

CIFAR-10 데이터셋을 torchvision을 사용하여 불러오고, 데이터 로더를 설정합니다.

from torchvision import datasets, transforms from torch.utils.data import DataLoader train_transform = transforms.Compose([ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), ]) test_transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), ]) train_dataset = datasets.CIFAR10(root="./data", train=True, download=True, transform=train_transform) test_dataset = datasets.CIFAR10(root="./data", train=False, download=True, transform=test_transform) train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True, num_workers=2) test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False, num_workers=2)

4단계: 모델 불러오기 및 초기 설정

사전 학습된 ResNet18 모델을 불러오고, 필요한 설정들을 정의합니다. 예시 코드에서는 별도의 ResNet18 구현 파일(resnet.py)이 필요하며, 사전 학습된 가중치 파일(resnet18.pth)을 로드합니다. 모델 구현 및 사전 학습 과정은 pytorch-cifar 레포지토리를 참고해주세요.

import torch from resnet import resnet18 # 사용자 정의 ResNet18 구현 파일 device = "cuda" if torch.cuda.is_available() else "cpu" model = resnet18(num_classes=10) checkpoint = torch.load("resnet18.pth") # 사전 학습된 가중치 파일 model.load_state_dict(checkpoint) model.to(device) import torch.nn as nn import torch.optim as optim criterion = nn.CrossEntropyLoss().to(device) optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4) # 조정된 learning rate scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=20) # 조정된 epoch

QAT는 가중치를 미세 조정하는 과정이므로, 일반 학습보다 낮은 learning rate와 적은 에폭 수를 사용하는 것이 일반적입니다.

5단계: QAT 수행 및 평가

학습 및 평가 함수를 정의하고, QAT를 수행합니다.

def train(num_epochs, model, train_loader, test_loader, optimizer, scheduler, criterion, device): for epoch in range(num_epochs): model.train() running_loss, total, correct = 0.0, 0, 0 for images, labels in train_loader: images, labels = images.to(device), labels.to(device) optimizer.zero_grad() outputs = model(images) loss = criterion(outputs, labels) loss.backward() optimizer.step() running_loss += loss.item() * images.size(0) _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() scheduler.step() train_acc = 100.0 * correct / total print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(train_loader.dataset):.4f}, Acc: {train_acc:.2f}%") test_acc = evaluate(model, test_loader, device) print(f"Test Accuracy: {test_acc:.2f}%") def evaluate(model, data_loader, device): model.eval() correct, total = 0, 0 with torch.inference_mode(): for images, labels in data_loader: images, labels = images.to(device), labels.to(device) outputs = model(images) _, predicted = torch.max(outputs, dim=1) total += labels.size(0) correct += (predicted == labels).sum().item() return 100.0 * correct / total train(num_epochs=20, model=model, train_loader=train_loader, test_loader=test_loader, optimizer=optimizer, scheduler=scheduler, criterion=criterion, device=device) # num_epochs 조정 필요 torch.save(model.state_dict(), "resnet18_qat.pth") # QAT 모델 저장 print("QAT 모델 저장 완료: resnet18_qat.pth")

6단계: ONNX 변환을 위한 준비

ONNX 변환을 위해 pytorch-quantization을 다시 초기화하고, QAT 모델을 불러옵니다. 이때 use_fb_fake_quant를 True로 설정하여 ONNX export 시 fake quantization 노드가 포함되도록 합니다.

from pytorch_quantization import quant_modules from pytorch_quantization import nn as quant_nn quant_nn.TensorQuantizer.use_fb_fake_quant = True # Fake quantization 노드 포함 quant_modules.initialize() device = "cuda" if torch.cuda.is_available() else "cpu" model = resnet18(num_classes=10) model.to(device) checkpoint = torch.load("resnet18_qat.pth") # QAT 모델 로드 model.load_state_dict(checkpoint) model.eval()

7단계: ONNX 변환

torch.onnx.export를 사용하여 모델을 ONNX 형식으로 변환합니다.

def export_to_onnx(model, onnx_filename, device): dummy_input = torch.randn(1, 3, 32, 32).to(device) # Dummy input torch.onnx.export( model, dummy_input, onnx_filename, input_names=["input"], output_names=["output"], opset_version=13, do_constant_folding=True ) print(f"ONNX 모델 저장 완료: {onnx_filename}") export_to_onnx(model, "resnet18_qat.onnx", device)

8단계: ONNX 모델 시각화

변환된 ONNX 모델을 Netron으로 시각화하면 Q/DQ 노드가 추가된 것을 확인할 수 있습니다. Netron은 신경망 모델을 시각적으로 표현해주는 도구입니다.

9단계: TensorRT 엔진 생성

ONNX 모델을 TensorRT 엔진으로 변환하여 NVIDIA GPU에서 더 빠른 추론 성능을 얻을 수 있습니다. TensorRT 엔진 생성에는 별도의 과정이 필요하며, 자세한 내용은 NVIDIA TensorRT Documentation을 참고하시기 바랍니다. 일반적으로 trtexec 도구를 사용합니다.

trtexec --onnx=resnet18_qat.onnx --int8 --saveEngine=resnet18_qat.engine

trtexec는 ONNX 모델을 NVIDIA GPU에서 실행하기 위한 최적화된 TensorRT 엔진으로 변환하는 도구입니다. 이렇게 생성된 TensorRT 엔진은 원본 모델보다 빠른 속도로 추론을 수행할 수 있습니다.

성능 비교 분석

이제 각 모델의 성능을 비교 분석해 보겠습니다. 먼저, CIFAR-10 데이터셋과 ResNet-18 모델을 사용한 실험 결과를 살펴보죠. FP32 모델, QAT를 적용하지 않은 INT8 TensorRT 엔진, PTQ를 적용한 INT8 TensorRT 엔진, QAT를 적용한 INT8 TensorRT 엔진의 성능을 아래 표에 정리했습니다.

표에서 볼 수 있듯이, 네 가지 모델 모두 CIFAR-10 이미지 분류에서 95% 이상의 정확도를 보여줍니다. CIFAR-10 데이터셋과 ResNet-18 모델은 비교적 작은 규모이기 때문에, FP32에서 INT8로 변환하더라도 성능 저하가 크지 않고, 심지어 PTQ의 성능이 QAT보다 잘 나오는 것을 확인할 수 있습니다.

하지만 앞서 언급했던 3D 객체 검출 예시에서는 다른 결과가 나타났습니다. 3D 객체 검출는 모델 구조가 복잡하고 데이터셋의 규모가 크기 때문에 양자화의 영향을 더 크게 받습니다.

3D 객체 검출예시에서는 PTQ를 적용했을 때 FP32 대비 약 3.99%의 성능 저하가 발생했지만, QAT를 적용했을 때는 0.86%의 성능 저하만 발생했습니다. 이처럼 모델 구조가 복잡하고 데이터셋 규모가 클수록 QAT의 효과가 더욱 뚜렷하게 나타나는 것을 알 수 있습니다. QAT는 훈련 과정에서 양자화 연산을 미리 고려하기 때문에 PTQ보다 양자화로 인한 성능 저하를 효과적으로 줄일 수 있습니다.

결론적으로, QAT는 PTQ에 비해 성능 저하를 최소화하면서 모델의 크기를 줄이고 연산 속도를 높일 수 있는 효과적인 양자화 기법입니다. 특히 복잡한 모델이나 대규모 데이터셋을 사용하는 경우 QAT를 적용하는 것이 더욱 유리합니다.

결론

이번 글에서는 로봇 ML 모델 경량화를 위한 훈련 후 양자화(PTQ)의 한계점과 이를 극복하기 위한 양자화 인식 훈련(QAT)에 대해 알아보았습니다. QAT는 훈련 과정에서 양자화를 고려하여 모델을 최적화하는 기법으로, PTQ보다 성능 저하를 최소화하면서 모델 크기를 줄이고 연산 속도를 높일 수 있습니다. 특히, 복잡한 모델이나 대규모 데이터셋을 사용하는 경우 QAT의 효과가 더욱 뚜렷하게 나타납니다.

하지만 QAT는 기존 학습 모델을 재사용할 수 없고 처음부터 학습을 다시 진행해야 한다는 단점이 있습니다. 또한 모델이나 데이터의 사이즈가 작을 경우 효과가 크지 않을 수 있습니다. 따라서, QAT 적용 시에는 추가적인 학습 시간과 자원 소모를 고려해야 합니다.

결론적으로, 로봇 ML 모델 경량화를 위해서는 PTQ와 QAT의 장단점을 비교하여 상황에 맞는 적절한 기법을 선택하는 것이 중요합니다. 비교적 복잡한 모델을 다룰 때, 그 정확도를 최대한 유지하면서 경량화를 달성해야 하는 경우, QAT는 효과적인 선택이 될 수 있습니다.

더 나아가, QAT의 성능을 더욱 향상시키기 위한 다양한 연구들이 진행되고 있습니다. 예를 들어, 양자화된 모델의 정확도를 높이기 위한 새로운 학습 방법이나, 양자화 과정에서 발생하는 정보 손실을 최소화하기 위한 기법들이 연구되고 있습니다. 앞으로도 이러한 연구들을 통해 로봇 ML 모델의 경량화 기술은 더욱 발전할 것으로 기대됩니다.

👉 관련 글 보러 가기

문종식