[ultralytics] Ultralytics YOLO에 INT8 Quantization-Aware Training(QAT) 도입하기

PR 링크: ultralytics/ultralytics#26083 상태: Merged | 변경: +258 / -28

들어가며

딥러닝 모델을 엣지 디바이스에 배포할 때, INT8 양자화는 성능 향상을 위한 필수적인 과정입니다. 하지만 일반적으로 사용되는 PTQ(Post-Training Quantization)는 학습이 끝난 후 양자화를 수행하기 때문에 필연적으로 정확도 손실이 발생합니다. 이번 Ultralytics의 업데이트에서는 quantize=8 옵션을 통해 학습 루프 내에서 가중치를 INT8 환경에 적응시키는 Quantization-Aware Training(QAT)을 도입하여, 배포 시 정확도 저하를 획기적으로 줄였습니다.

코드 분석

1. engine/trainer.py: QAT 상태 관리 및 복원

학습 중 QAT를 적용하기 위해 모델을 준비하고, 체크포인트 재개 시 양자화 상태를 복원하는 로직이 추가되었습니다. 특히 restore_qat 함수는 체크포인트에 저장된 양자화 범위를 모델에 다시 적용합니다.

# Before: 체크포인트 로드 시 양자화 상태 복원 로직 부재
# After: 체크포인트 로드 후 QAT 상태 복원
if is_qat(model):
    restore_qat(model, ckpt["modelopt"])

2. engine/exporter.py: QAT 모델 전용 내보내기

QAT 모델은 이미 학습 과정에서 양자화 범위가 고정되어 있으므로, 내보내기 단계에서 별도의 보정(Calibration) 과정이 필요 없습니다. 이를 위해 is_qat를 확인하여 불필요한 과정을 생략합니다.

# Before: 모든 모델에 대해 보정 데이터 요구
# After: QAT 모델은 보정 단계 생략 및 Q/DQ 노드 생성
self.qat = is_qat(model)
if self.qat:
    assert fmt in {"onnx", "engine"}
    self.args.quantize = 8

왜 이게 좋은가

QAT는 단순히 양자화를 수행하는 것이 아니라, 모델이 INT8 환경에서 최적의 성능을 내도록 가중치를 미세 조정(Fine-tuning)합니다. 실제 RTX PRO 6000 환경에서 yolo26s 모델을 테스트한 결과, PTQ 대비 정확도 손실을 약 3%p 이상 개선하는 성과를 보였습니다.

• 정확도 유지: PTQ에서 발생하는 급격한 정확도 하락을 방지합니다.
• 배포 편의성: 학습된 체크포인트에 양자화 범위가 포함되어 있어, 별도의 보정 데이터셋 없이 즉시 onnx나 engine으로 내보낼 수 있습니다.
• 교훈: 모델의 크기가 커질수록 양자화로 인한 손실이 커지므로, QAT는 특히 대형 모델이나 정밀도가 중요한 환경에서 필수적인 최적화 전략입니다.

결론

이번 업데이트를 통해 Ultralytics는 사용자가 복잡한 양자화 지식 없이도 quantize=8 옵션 하나만으로 고성능 INT8 모델을 확보할 수 있게 했습니다. torch.compile과의 호환성 제약은 존재하지만, 배포 성능 측면에서 얻는 이득이 훨씬 큽니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.compile.html
• https://github.com/NVIDIA/TensorRT-Model-Optimizer

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.