[논문 리뷰] CompressAI: a PyTorch library and evaluation platform for end-to-end compression research

Jean Bégaint, Fabien Racapé외 2명|arXiv (Cornell University)|2020. 11. 05.

Advanced Data Compression Techniques참고 문헌 22인용 수 235

한 줄 요약

CompressAI는 PyTorch 기반 라이브러리, 사전 학습된 모델, 연구, 학습 및 엔드-투-엔드 이미지 압축 코덱의 전통 코덱에 대한 벤치마킹 도구를 제공합니다.

ABSTRACT

This paper presents CompressAI, a platform that provides custom operations, layers, models and tools to research, develop and evaluate end-to-end image and video compression codecs. In particular, CompressAI includes pre-trained models and evaluation tools to compare learned methods with traditional codecs. Multiple models from the state-of-the-art on learned end-to-end compression have thus been reimplemented in PyTorch and trained from scratch. We also report objective comparison results using PSNR and MS-SSIM metrics vs. bit-rate, using the Kodak image dataset as test set. Although this framework currently implements models for still-picture compression, it is intended to be soon extended to the video compression domain.

연구 동기 및 목표

• 엔드-투-엔드 이미지(향후 비디오) 압축 코덱을 연구하고 개발하기 위한 PyTorch 기반 플랫폼 제공.
• 도메인 특화 구성요소로 최첨단 학습 이미지 압축 모델 재구현.
• 게시된 결과를 재현하고 확장하기 위한 훈련 가중치 및 모델 집(모델 zoo) 제공.
• PSNR 및 MS-SSIM을 사용하여 학습된 코덱과 전통 코덱 간의 객관적 벤치마킹 용이화.
• 학습에서 성능 보고까지 엔드-투-엔드 학습 및 평가 파이프라인 활성화.

제안 방법

• 압축을 위한 도메인 특화 PyTorch 레이어 및 연산 구현(엔트로피 모델, 양자화, 색상 변환).
• PyTorch에서 최첨단 학습 이미지 압축 모델 재구현 및 사전 학습 가중치 제공.
• 다양한 품질/비트레이트 설정에 대해 다운로드 가능한 사전 학습 네트워크를 갖춘 모델 zoo 제공.
• 손실 함수가 왜곡(distorion)과 비트레이트 항을 결합한 학습 및 평가 파이프라인 제공.
• 데이터세트에서 모델 평가 및 전통 코덱과의 벤치마킹을 위한 유틸리티와 CLI 제공.
• Kodak/CLIC 데이터세트에서 성능 벤치마크를 수행하고 비트레이트별 PSNR 및 MS-SSIM 보고.

실험 결과

연구 질문

• RQ1 CompressAI에서 재구현된 학습 이미지 압축 모델이 발표된 Rate-Distortion 결과를 얼마나 가까이 재현할 수 있는가?
• RQ2 다양한 비트레이트에서 학습된 엔드-투-엔드 코덱은 JPEG, JPEG2000, WebP, HEVC, AV1, VVC 같은 전통 코덱과 PSNR 및 MS-SSIM에서 어떻게 비교되는가?
• RQ3 PyTorch 기반 도구가 학습, 평가 및 학습된 코덱과 기존 표준 간의 비교를 간소화할 수 있는가?
• RQ4 이미지에서 비디오로 엔드-투-엔드 압축 도구를 확장하는 타당성과 경로는 무엇인가?
• RQ5 CompressAI 내에서 어떤 학습 구성(손실, 비트 깊이, 데이터세트)이 최신 성능을 달성하는가?

주요 결과

• 재구현된 모델(요인화된 프라이어, 하이퍼프라이어, 가우시안 혼합, 자기회귀/하이퍼프라이어 조합)은 원작에서 훈련 시 초기값으로 거의 비슷한 결과를 재현한다.
• 사전 학습 가중치를 사용하여 보고된 결과를 가까이 재현하고 다양한 지표나 비트레이트에 대해 빠르게 미세 조정 가능.
• 학습된 코덱은 Kodak의 PSNR 및 MS-SSIM 벤치마크에서 전통 방법(JPEG, JPEG2000, WebP)에 비해 크게 우수하며, HEVC, AV1과 현대 표준과 경쟁력이 있으며 저비트레이트에서 VVC에 근접.
• 플랫폼은 엔드-투-엔드 학습 코덱이 PyTorch 기반의 통합 워크플로우 내에서 학습 및 평가될 수 있음을 시연하고 표준 데이터세트에서 재현성을 보인다.
• CompressAI의 공개 채택은 산업계 및 학술 연구에서의 사용과 표준화 그룹 내 DNN 기반 비디오 코딩 논의에 영향을 미친 것으로 입증된다.