OwLite로 YOLOv5 모델 경량화하기 1/3 — OwLite 환경 세팅, Baseline 등록
OwLite 도입의 첫 단계로, 사용 환경을 설정하고 Baseline을 등록하는 과정을 소개합니다.
AI는 점점 더 많은 분야에서 핵심적인 역할을 수행하고 있습니다. 컴퓨터 비전, 자연어 처리, 음성 인식 등 다양한 분야에서 AI의 활용이 증가하고 있으며, 많은 기업들이 자체 AI 모델을 개발하고 이를 서비스에 도입하기 위해 노력하고 있습니다. 그러나 AI 모델의 성능이 향상되고 기능이 다양해질수록, 모델의 크기 또한 커져 기업들이 실질적인 배포와 서비스에 어려움을 겪고 있습니다.
이러한 상황에서 스퀴즈비츠는 AI 모델 경량화 도입을 돕는 AI 경량화 툴킷 OwLite를 출시했습니다. OwLite는 기존의 Python 환경에서 최소한의 코드 변경만으로 빠르게 경량화를 적용할 수 있게 설계되었습니다. 처음 경량화를 시도하는 분들도 손쉽게 사용할 수 있는 UX를 고려하였고, 튜토리얼을 포함한 사용문서, 실제 적용 예시를 제공하고 커뮤니티 등의 적극적인 지원으로 사용의 어려움을 최소화하기 위해 노력하고 있습니다:)
이번 글에서는 가상의 고객이 OwLite로 AI 모델을 경량화하는 과정을 통해 OwLite를 실제 모델에 적용하는 방법과 이를 통해 얻을 수 있는 성과를 보여드리고자 합니다. OwLite를 실제 작업에 적용하면서 마주할 수 있는 다양한 상황들을 함께 살펴보시죠!
—
“Edge device에서 Object Detection AI를 사용해야 하는데, YoloV5 모델의 inference에 걸리는 시간이 너무 길어 실시간 서비스가 어렵습니다.”
고객은 스마트 CCTV 기술을 개발하는 한 스타트업의 엔지니어로, 아래와 같은 상세 스펙을 보유하고 있습니다.
Model: YOLOv5-s
Device: NVIDIA Jetson AGX Orin
Input shape: -1 × 3 × 640 × 640 (Dynamic batch size)
Current accuracy: 33.3 mAP@.5 : .95
Current latency: 16.088 ms
Dataset: coco 2017
이러한 사례를 바탕으로 모델의 성능은 유지하면서 inference 속도를 개선하기 위해 OwLite로 YOLOv5 모델을 경량화 해보겠습니다.
계정 생성 및 로그인
OwLite 서비스를 사용하려면 먼저 OwLite 웹페이지에서 계정을 생성해야 합니다. 계정을 생성한 후 로그인하면 나타나는 홈 화면에서 현재 구독 중인 Plan 정보와 최근에 등록한 Experiments, 그리고 웹 사용 튜토리얼을 확인할 수 있습니다.
OwLite python package 설치
pip install owlite --extra-index-url https://pypi.squeezebits.com/OwLite의 코드 베이스는 Python 형태로 제공됩니다. 모델 개발 환경에서 터미널에 위 명령어를 입력함으로써 최신 버전의 패키지를 설치할 수 있습니다.
OwLite 코드 실행
모델을 경량화하기 위해 기존에 사용하던 개발 환경에 OwLite 코드를 추가합니다. 이번 글에서는 YOLOv5 모델을 학습시키는 train.py 코드를 기반으로 OwLite 코드를 작성했습니다.
1. CLI 설정
코드를 실행하기 전에, 개발 환경에서 OwLite에 로그인합니다. OwLite의 CLI를 통해 사용할 계정과 디바이스를 설정할 수 있습니다.
~ owlite login
Enter your email: <email>
Enter your password: <password>
OwLite [INFO] Logged in as sarangkim
OwLite [INFO] Your price plan: BUSINESS
OwLite [INFO] Your workgroup: sqzb
OwLite [INFO] Your authentication token is saved at ~/.cache/owlite/tokens로그인 후, CLI를 통해 계정 관리자가 설정해 둔 디바이스에 연결할 수 있습니다. OwLite는 이때 연결된 디바이스에서 TensorRT 엔진을 생성하고 속도를 측정하므로, 경량화한 모델의 속도를 측정할 디바이스를 선택해야 합니다. 이번 글에서는 고객이 대상으로 하는 디바이스인 NVIDIA Jetson AGX Orin에 연결하였습니다.
~ owlite device connect
OwLite [INFO] Available devices:
0: NVIDIA RTX A6000
1: NVIDIA Jetson AGX Orin
Enter the index of the device you want to connect to: 1
OwLite [INFO] Connected to the device 'NVIDIA Jetson AGX Orin' at 'NEST' (https://nest.owlite.ai)2. OwLite 패키지 불러오기
이제 본격적으로 작업하던 코드에 OwLite 코드를 추가해보겠습니다. OwLite 파이썬 패키지 사용을 위해, 파이썬 코드의 상단에서 OwLite 패키지를 불러옵니다.
import owlite3. Baseline 등록
OwLite로 모델을 경량화하는 과정은 Baseline 등록-> 경량화 옵션 설정 -> 모델에 경량화 옵션 적용 -> 경량화된 모델의 Benchmark 측정 순서로 이루어집니다.
먼저, 파이썬 코드 내에 Baseline 등록을 위한 코드를 추가합니다.
1) Init (Project, Baseline)
owl = owlite.init("object_detection", "yolov5s")코드의 상단에서 OwLite의 Init 함수를 사용해 Project(object_detection) 와Baseline(yolov5s)을 생성합니다. 생성한 인스턴스는 변수 owl에 저장됩니다.
2) Convert (model)
model = owl.convert(model.eval(), example_input)다음으로, YOLOv5-s 모델을 OwLite가 사용할 수 있는 형식으로 변환하기 위해 Convert 함수를 추가합니다.
이때, Convert 함수를 통해 모델을 변환하기 전, 모델 선언, checkpoint 로드 등 모델을 정의하는 데 필요한 작업을 완료해야합니다. 단, DataParallel은 모델 변환 후에 적용합니다.
또한, Convert 함수에는 모델 입력의 예시를 함께 넣어야 합니다. 이 입력 예시는 추후 모델을 forward할 때 사용할 입력과 동일한 형식이어야 합니다.
이 코드에서는 아래 두 가지 조건을 고려하여 example_input을 생성했습니다:
- dataloader의 이미지 사이즈
- DataParallel 사용을 고려한 배치 사이즈
3) Validate, Log (mAP)
owl.log(mAP=results[3])OwLite의 Log 함수를 사용해 경량화 과정에서 만들어지는 각 모델의 정보를 기록하고 비교할 수 있습니다. 이 코드에서는 YOLOv5 모델의 validation 결과 중 mAP 값을 기록합니다.
이 과정에서, 학습을 진행하지 않고도 validation만 수행할 수 있도록 기존 iteration 코드를 수정하고, mAP를 기록하는 코드를 추가했습니다.
4) Export, Benchmark
owl.export(de_parallel(model))
owl.benchmark()마지막으로 연결된 디바이스에서 TensorRT 엔진을 생성하고 latency를 측정하기 위해, 모델을 Export한 후 Benchmark를 실행합니다.
OwLite의 Benchmark는 Export 함수로 생성된 모델로 수행됩니다. 따라서 실제 배포 환경과 가장 유사한 결과를 얻기 위해, 학습이 완료되고 모델의 파라미터가 최종적으로 업데이트 된 후에 Export 함수를 실행합니다.
5) 각종 에러 상황 해결
앞의 1)~4)에서 추가한 코드를 실행할 때, 환경에 따라 몇 가지 에러가 발생할 수 있습니다. 이러한 경우 owlite-examples를 참고하여 발생하는 에러를 해결할 수 있습니다.
만약 사용 중인 모델이 예시에서 다루어지지 않았거나, 예기치 않은 문제가 발생했다면 Help Center를 방문해주세요! Community에 질문을 남기면 스퀴즈비츠의 엔지니어와 다른 사용자들에게 도움을 받으실 수 있습니다:)
이번 글의 YOLOv5는 owlite-examples에 있는 모델이기 때문에 예시의 apply_owlite.patch 파일을 참고해서 문제를 해결해보겠습니다.
Error #1 : model convert error
모델의 forward 함수에 있는 warnings.catch_warnings() 블록이 Convert 과정에서 에러를 발생시키고 있었습니다.
해당 블록을 제거하니 모델 변환 코드가 문제없이 동작하는 것을 확인했습니다.
이후 validate 과정에서 발생하는 에러도 함께 해결해 보겠습니다.
Error #2 : deepcopy error
ModelEMA 부분이 해당 에러를 발생시키는 것으로 보입니다.
이 코드에서는 EMA를 사용하지 않기 때문에 해당 코드를 제거하였습니다.
Error #3: Invalid input size error
이번에는 vallidate를 위한 forward 과정에서 에러가 발생했네요. Input shape가 맞지 않다는 에러입니다. 문서의 Trobuleshooting 섹션에 기록된 대로, OwLite로 변환한 모델은 example input과 동일한 shape의 데이터만 입력 할 수 있습니다.
이 문제를 해결하기 위해, dataloader의 batch_size를 Convert 함수에 넣어준 example_input과 동일한 batch_size로 변경하고, 모든 데이터가 같은 batch_size를 갖도록 drop_last 인자를 True로 설정합니다.
또한, YOLOv5 코드는 데이터셋의 이미지 비율만 같고 크기가 달라질 수 있는 Dataloader를 생성하기 때문에, 모든 이미지가 동일한 크기를 가지도록 Dataloader 생성 코드를 수정했습니다.
6) 등록된 Baseline 확인
프로젝트 리스트 화면에서 Baseline 모델과 Log 함수로 기록한 mAP가 모두 잘 등록되었음을 확인할 수 있습니다. 또한, Jetson AGX Orin 디바이스에서 측정된 TensorRT 엔진의 latency 값도 기록되었습니다. OwLite는 Baseline의 TensorRT 엔진을 빌드할 때 FP16 precision 모드를 활성화하여 빌드합니다.
—
다음 편에서는 등록한 Baseline 모델에 다양한 경량화 옵션을 적용하고, latency를 개선하는 과정을 살펴보겠습니다.
2편은 아래에서 이어집니다.
지금 바로 OwLite를 사용해 보고 싶다면 여기를 방문하세요!
