[論文レビュー] CompressAI: a PyTorch library and evaluation platform for end-to-end compression research

Jean Bégaint, Fabien Racapé+2 more|arXiv (Cornell University)|Nov 5, 2020

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 を提供する。
• 歪みとレートの項を組み合わせた損失関数を含む訓練・評価パイプラインを提供する。
• データセット上でモデルを評価し、従来のコーデックと比較するためのユーティリティと CLI を提供する。
• Kodak/CLICデータセットでの性能をベンチマークし、ビットレートごとにPSNRとMS-SSIMを報告する。

実験結果

リサーチクエスチョン

• RQ1 CompressAI で再実装された学習済み画像圧縮モデルは、公開されたレート歪み結果を再現できるか？
• RQ2 さまざまなビットレートで、学習済みエンドツーエンドコーダは従来のコーダ（JPEG、JPEG2000、WebP、HEVC、AV1、VVC）とPSNRおよびMS-SSIMでどう比較されるか？
• RQ3 PyTorchベースのツールキットは、学習済みコーデックの訓練、評価、比較を合理化できるか？
• RQ4 画像から映像へエンドツーエンド圧縮ツールを拡張する際の実現性と動向は？
• RQ5 どの訓練設定（損失、ビット深度、データセット）が CompressAI 内で最先端の性能を生み出すか？

主な発見

• 再現実装モデル（factorized prior、hyperprior、Gaussian mixture、autoregressive/hyperprior の組み合わせ）は、ゼロから訓練した場合、元の論文で報告された結果に近い結果を再現する。
• 事前訓練済み重みは、報告された結果をほぼ再現し、異なる指標やビットレートのための素早い微調整を可能にする。
• 学習済みコーダは、PSNRとMS-SSIMのベンチマークで従来手法（例: JPEG、JPEG2000、WebP）を大きく上回り、KodakでHEVC、AV1、低ビットレートでVVCに近づく競争力を示す。
• このプラットフォームは、エンドツーエンド学習コーダを統一された PyTorch ベースのワークフロー内で訓練・評価できることを示し、標準データセットで再現性を実証している。
• CompressAI の公開の採用は、産業界や学術研究での利用と、標準化グループ内の DNN ベースの映像圧縮議論に影響を与えていることによって示されている。