OwLite로 YOLOv5 모델 경량화하기 3/3 — QAT, Dynamic batch size 엔진 생성
마지막으로, OwLite를 통해 경량화 모델을 학습시키는 QAT와 Dynamic batch size 엔진을 생성하는 과정을 소개합니다.
스퀴즈비츠에서 출시한 AI 경량화 툴킷인 OwLite를 활용하면, 간단한 코드 수정과 직관적인 에디터로 AI 모델을 쉽게 경량화 할 수 있습니다. 이번 마지막 편에서는 OwLite를 사용해 YOLOv5 모델을 경량화하는 방법을 마무리하며, 추가 학습을 통해 mAP를 Baseline 모델 수준으로 향상시키고, dynamic batch size 모델을 생성하는 과정을 소개해드리겠습니다.
QAT (Quantization Aware Training)
QAT(Quantization Aware Training)는 양자화를 고려하여 학습하는 방법으로, PTQ(Post-Training Quantization) 적용 시 발생할 수 있는 정확도 손실을 다시 복구하기 위해 사용됩니다. 이를 통해 단순히 PTQ를 적용할 때보다 더 높은 정확도를 유지하면서도 모델을 경량화 할 수 있는 장점이 있습니다.
앞선 2편에서는 YOLOv5 모델에 PTQ를 적용하여 32.5 mAP를 달성했습니다. 하지만 Baseline 모델의 mAP가 33.3이었던 것을 고려하면, 양자화 후 꽤 큰 성능 저하가 발생했습니다. 그래서 이번 편에서는 PTQ 후 추가로 QAT를 적용하여 mAP를 더욱 향상시키는 방법을 살펴보려고 합니다.
OwLite를 이용하면 PTQ를 통해 경량화 된 모델을 Fine-Tuning하면서 QAT를 적용해 성능을 최적화할 수 있습니다. QAT는 학습 중에 양자화 오차를 반영하여 파라미터를 조정하기 때문에, PTQ만 적용한 모델보다 더 나은 성능을 기대할 수 있습니다.
QAT도 이전과 동일하게 train.py 파일을 사용하여 진행했으며, 모델 배포 환경인 Jetson Orin AGX이 아닌 개발 환경인 NVIDIA A6000에서 실행했습니다. 이렇게 OwLite를 사용할 때에는 학습/개발을 진행하는 서버와 배포/벤치마크하는 환경을 다르게 사용해도 무방합니다.
owl = owlite.init(project="object_detection", baseline="yolov5s", experiment="recommend_mse_qat", duplicate_from="recommend_mse")PTQ가 완료된 후 QAT를 실시하기 때문에, PTQ를 적용한 모델 중 가장 높은 mAP를 달성한 recommend_mse Experiment를 복제하여 recommend_mse_qat라는 새로운 Experiment를 생성하도록 Init 함수를 수정했습니다.
학습에 필요한 하이퍼파라미터는 owlite-examples와 동일하게 설정한 뒤 코드를 실행했습니다. 그런데, 모델 학습 과정에서 에러가 발생했네요.
Error : Model output Index Error
발생한 에러 메시지는 index 3 is out of bounds for dimension 0 with size 3으로, compute_loss 함수에 전달된 output의 index가 잘못되었음을 나타냅니다. compute_loss 함수는 모델의 예측 결과와 실제 라벨을 비교하여 손실(loss)을 계산하는데, 이 과정에서 잘못된 index로 인해 에러가 발생한 것입니다. 이 에러는 모델이 eval 모드로 전환되면서 기존 torch 모델과 output 구조가 달라져서 발생했습니다.
YOLOv5 모델의 forward 함수는 다음과 같습니다:
YOLOv5 모델은 train 모드와 eval 모드에서 output의 형식이 다릅니다. train 모드에서는 output이 단일 tensor인 반면, eval 모드에서는 output이 tuple로 구성되며, 두 번째 요소에 x가 포함됩니다. OwLite가 Convert 함수에서 모델을 eval 모드로 변환했기 때문에, 이 부분을 고려해야 합니다.
이를 해결하기 위해 compute_loss 함수에 전달되는 pred가 OwLite로 변환된 모델의 output에서 두 번째 요소를 사용하도록 코드를 수정했습니다. 이렇게 compute_loss 함수의 input을 기존 torch 모델이 예상하는 output 구조와 일치시키면서 문제를 해결할 수 있었습니다.
owlite-example을 따라 3 epochs 동안 QAT를 진행했고, Baseline 모델과 거의 동일한 수준의 mAP를 달성할 수 있었습니다. (Baseline: 33.3 mAP, Quantized+QAT: 33.1 mAP)
Insights
OwLite는 여러 Experiment의 결과를 한눈에 비교할 수 있는 기능을 제공합니다. 프로젝트 페이지의 Insights 탭에서 Benchmark 결과를 다각형 그래프로 확인할 수 있고, Latency, Model size, peak VRAM usage뿐만 아니라 사용자가 정의한 각종 Log 값들 (여기서는 mAP)도 시각적으로 비교할 수 있습니다.
Insights 탭을 통해 지금까지 수행한 실험들을 선택하여 Latency, mAP, Model size, peak VRAM을 한눈에 비교할 수 있었습니다. 그 결과, recommend_mse_qat Experiment가 가장 빠르고 작으면서도 성능이 유지된 모델임을 확인했습니다.
이 모델은 다음과 같은 경량화 과정을 거쳤습니다:
1. Recommend 경량화 옵션 적용: OwLite의 Recommend Setting을 통해 경량화 옵션을 적용하여 Model size와 Latency를 줄였습니다.
2. MSE Calibrator로 PTQ: MSE Calibrator를 사용하여 Post-Training Quantization(PTQ)을 수행해 모델을 경량화했습니다.
3. QAT로 mAP 회복: Quantization Aware Training(QAT)과 Fine-Tuning을 통해 경량화된 모델의 mAP를 Baseline 모델 수준으로 회복했습니다.
recommend_mse_qat 모델은 Latency, Model size, 그리고 peak VRAM 측면에서 모두 우수한 성능을 보였음을 확인할 수 있었습니다.
이렇게 OwLite를 활용하여 YOLOv5 모델의 경량화 및 최적화를 성공적으로 마칠 수 있었습니다.
Dynamic Batch Size Model
프로젝트 리스트의 Input Signature를 확인해보면, 등록한 YOLOv5 모델의 입력 사이즈가 8x3x640x640으로, 배치사이즈가 8로 고정된 입력을 받고 있음을 알 수 있습니다. 이러한 Experiment에서는 고정된 input shape를 처리하는 TensorRT 모델이 생성됩니다.
그러나 1편에서 정의한 고객은 dynamic batch size의 모델을 필요로 하고 있습니다. TensorRT는 기본적으로 dynamic batch size 기능을 지원하며, OwLite 역시 dynamic batch size의 TensorRT 모델을 생성할 수 있도록 dynamic batch size 옵션을 제공합니다.
1. Experiment 복제 및 모델 로드
우선, dynamic batch size 모델을 만들기 위한 새로운 Experiment를 생성합니다.
owl = owlite.init(project="object_detection", baseline="yolov5s", experiment="recommend_dynamic", duplicate_from="recommend_mse_qat")앞의 모델의 경량화 옵션을 그대로 적용하기 위해 recommend_mse_qat Experiment를 복제해서 recommend_dynamic Experiment를 생성하고, QAT로 얻은 checkpoint를 불러옵니다. 이후, Export와 Benchmark의 옵션을 설정함으로써 dynamic batch size의 모델을 생성할 수 있습니다.
1) 모델의 input signature 확인
Dynamic batch size 옵션에 모델에서 dynamic batch size를 적용할 input 이름을 사용하기 위해, 프로젝트 리스트 탭의 INPUT SIGNATURE을 확인합니다. 등록된 YOLOv5s Baseline 모델의 input 이름이 “x”임을 확인했습니다. 이 이름을 기준으로 dynamic batch size를 설정할 input을 선택합니다.
2) Export의 dynamic_axis_options 설정
owl.export(de_parallel(owlite_model),
dynamic_axis_options={
"x":{
"axis": 0
}
}
)Export 시 dynamic_axis_options에서 input의 batch size를 의미하는 축을 설정합니다. 앞에서 확인한 모델의 input 이름이 “x”였고, 0번째 축이 batch size를 의미하므로 코드를 위와 같이 작성했습니다.
3) Benchmark의 dynamic_input_options 설정
TensorRT는 엔진을 생성할 때 주어진 input shape에 맞게 최적화를 수행하고, 지정된 입력에 대해 벤치마크를 실행합니다. 따라서 Export에서 dynamic batch size를 설정한 모델을 Benchmark하려면 모델을 TensorRT 엔진으로 변환할 때 사용할 구체적인 batch size 값을 지정해야 합니다.
owl.benchmark(
dynamic_input_options={
"x": {
"min": 1,
"opt": 32,
"max": 64,
"test": 32,
}
}
)Benchmark의 dynamic_input_options에 batch size의 최소(min), 최적화(opt), 최대(max) 범위를 설정하고, 테스트할 batch size (test)를 입력합니다. 이 값들은 TensorRT 엔진을 생성할 때 사용되며, 엔진은 이 범위 내에서 최적화됩니다. 여기서는 최소 1, 최적화 32, 최대 64, 테스트 시 batch size는 32로 설정했습니다.
4) 결과 확인
코드를 실행한 후, 프로젝트 리스트 탭에서 dynamic batch size 모델을 batch_size=32로 benchmark한 latency 결과를 확인할 수 있었습니다.
2. Baseline 모델의 dynamic batch_size 변환
Dynamic batch size의 경량화 모델을 생성했지만, Baseline 모델은 고정된 batch size로 설정되어 있었기 때문에 두 모델의 latency를 직접 비교하기 어려웠습니다. 따라서, Baseline 모델에도 동일한 dynamic batch size 옵션을 적용해 두 모델 간의 성능을 1:1로 비교할 수 있도록 했습니다.
owl = owlite.init(project="object_detection", baseline="yolov5s", experiment="dynamic_batch")Baseline과 같은 dynamic batch size를 적용하려면, 경량화 옵션이 없는 새로운 Experiment를 생성하면 됩니다. 여기서는 존재하지 않는 이름인 dynamic_batch 라는 Experiment를 생성했습니다. Export의 dynamic_axis_options Benchmark dynamic_input_options는 앞의 recommend_dynamic Experiment를 생성할 때와 동일한 값을 입력했습니다.
프로젝트 리스트 탭에서 동일한 input shape로 경량화 전/후 모델을 benchmark한 결과를 확인할 수 있습니다.
FP16 precision 모드로 빌드된 “dynamic_batch” 모델과 INT8 경량화가 적용된 “recommended_dynamic” 모델을 비교한 결과는 다음과 같습니다:
• 모델 사이즈: 약 64.1% 감소 (15.517 MB -> 8.365 MB)
• Latency: 약 27.2% 감소 (60.984 ms -> 44.419 ms)
• 모델 정확도(mAP): 약 0.2%p 하락 (33.3 mAP -> 33.1 mAP)
이로써 OwLite의 INT8 경량화를 통해 모델 사이즈와 Latency는 크게 개선되었고, 정확도에는 최소한의 영향만 미쳤다는 점을 확인할 수 있었습니다.
—
이렇게 OwLite를 활용하면 기존 모델 학습 환경에 간단한 코드만 추가함으로써 쉽게 경량화 옵션을 적용할 수 있습니다. 사용 과정에서 발생한 어려움은 문서, 적용 예시, 커뮤티니를 참고해 빠르게 문제를 해결할 수 있습니다.
또한, OwLite의 에디터를 통해 모델 구조를 직관적으로 파악할 수 있으며, 시각화된 그래프에서 옵션을 간편하게 조정할 수 있습니다. 마지막으로, Insights 탭에서는 여러 실험 결과를 비교해 최적의 모델을 선택하는 것도 가능합니다.
OwLite는 YOLO 모델뿐만 아니라 다양한 AI 모델의 경량화를 지원합니다. 지금, 여기를 방문해 여러분의 모델을 직접 경량화 해 보세요!
