ハン・ソン、チュー・ジュンヤンらがGAN圧縮方式を提案：計算電力消費は1/9以下、オープンソース化

GAN は大幅に圧縮することもできます。MIT ハン・ソン氏のチームによる最新の研究は、多くの研究者を興奮させています。

[[319638]]

生成モデルGANは、機械学習の分野で最も重要な開発方向の1つです。しかし、このようなアルゴリズムには膨大な計算能力が必要であり、ほとんどの研究者にとって新しい結果を導き出すのは困難です。近年、この分野は大規模な機関によって独占される傾向にあります。

しかし最近、マサチューセッツ工科大学 (MIT)、Adobe、上海交通大学の研究者らが、条件付き GAN を圧縮するための一般的な方法を提案しました。この新技術は、視覚的な忠実度を維持しながら、pix2pix、CycleGAN、GauGANなどの広く使用されている条件付きGANモデルの計算複雑性を1/9～1/21に削減します。この方法は、さまざまなジェネレータ アーキテクチャ、学習目標、ペア設定または非ペア設定に適用できます。

この研究論文はCVPR 2020カンファレンスに掲載されており、GAN圧縮フレームワークのPyTorchバージョンもオープンソース化されています。

プロジェクトリンク: https://github.com/mit-han-lab/gan-compression

圧縮後のGANの具体的なパフォーマンスはどれくらいですか?研究者らが披露したデモでは、CycleGAN を使って動画内の馬にシマウマの縞模様を追加するために必要な計算能力は 1/16 以下で、フレーム レートは 3 倍に増加し、効果も向上しました。

同研究所が使用するハードウェア プラットフォームは、NVIDIA のエッジ AI コンピューティング チップ Jetson Xavier GPU であることは特筆に値します。公式データによると、Jetson Xavier の INT8 コンピューティング能力は 22+10TOPS であるのに対し、Snapdragon 865 は 15TOPS です。圧縮されたGANはロボットやドローンなどの小型デバイスでも動作できるようになり、将来的には携帯電話に搭載できるようになるのも目前だそうです。

論文リンク: https://arxiv.org/pdf/2003.08936v1.pdf

概要

敵対的生成ネットワーク (GAN) は、非常にリアルな画像の合成に優れています。 GAN のバリエーションである条件付き生成的敵対的ネットワーク (cGAN) は、多くのコンピューター ビジョンおよびグラフィックス アプリケーションで制御可能な画像合成を可能にします。ただし、これらのアプリケーションのほとんどでは、人と対話するためのモデルが必要なので、より優れたユーザー エクスペリエンスを得るには、低遅延のデバイスが必要です。

しかし、最近導入された cGAN の中には、現在の認識畳み込みニューラル ネットワーク (CNN) よりも 1 ～ 2 桁も計算負荷が高いものもあります。たとえば、GanGAN はイメージごとに 281G Mac を消費しますが、MobileNet-v3 は 0.44G Mac しか必要としないため、前者をインタラクティブな展開に使用するのは困難です。

さらに、現時点では、ほとんどのエッジ デバイスはコンテンツやバッテリーなどのハードウェアによって制限されており、これもエッジ デバイスへの GAN の展開を妨げています。

そこで、画像合成分野における GAN と cGAN のこれらの問題に基づいて、Han Song 氏のチームは、GAN の推論時間と計算コストを削減するために使用される一般的な圧縮方法である GAN 圧縮を提案しました。同時に、生成モデルの圧縮には 2 つの基本的な困難が伴います。GAN のトレーニングは、特にペアになっていない場合には不安定であり、ジェネレーターは CNN とは異なるため、既存の CNN 設計を使用することは困難です。この問題に対処するために、チームは元の教師ジェネレーターの中間表現レイヤーから、対応する生徒ジェネレーターのレイヤーに知識を転送しました。

トレーニング コストを削減するために、チームは、すべての可能な数のチャネルを含む「一度限りのネットワーク」をトレーニングすることにより、モデル トレーニングとアーキテクチャ検索を分離しました。この「ワンス・フォー・オール・ネットワーク」は、重み共有を通じて多数のサブネットワークを生成することができ、各サブネットワークのパフォーマンスはトレーニングなしで評価できます。このモデルは、モデルアーキテクチャ、学習アルゴリズム、または監督設定 (ペアまたは非ペア) に関係なく、さまざまな条件の GAN モデルに適用できます。

研究チームは多数の実験を通じて、この方法により、生成される画像の忠実度を損なうことなく、広く使用されている 3 つの GAN モデル (pix2pix、CycleGAN、GauGAN) の計算複雑性を 1/9 ～ 1/21 に削減できることを証明しました。

具体的な方法

インタラクティブ アプリケーション用の条件付き生成モデルを圧縮することが、主に 2 つの理由から困難であることはよく知られています。まず、GAN の動的トレーニングは本質的に非常に不安定です。次に、認識モデルと生成モデル間のアーキテクチャ上の大きな違いにより、既存の CNN 圧縮アルゴリズムを直接使用することが困難になります。

これらの理由から、研究者は効率的な生成モデル専用のトレーニング スキームを提案し、ニューラル アーキテクチャ検索 (NAS) を使用して圧縮率をさらに高めました。 GAN 圧縮フレームワークの全体的なアーキテクチャを以下の図 3 に示します。ここでは、ResNet ジェネレータを例として使用しています。同じフレームワークをさまざまなジェネレータ アーキテクチャと学習目標に適用できることを強調することが重要です。

図 3: この論文の GAN 圧縮フレームワークの全体的なアーキテクチャ。

目的関数

1. ペア学習と非ペア学習の統合

トレーニングの目的が広範囲にわたるため、一般的な圧縮フレームワークを構築することは困難です。この問題に対処するために、教師モデルが元々どのようにトレーニングされたかに関係なく、モデル圧縮設定でペア学習と非ペア学習を統合します。元の教師ジェネレータ G′ が与えられた場合、ペアになっていないトレーニング設定をペアになっている設定に変換します。ペアになっていない設定では、元のジェネレーターの出力を真の値として扱い、ペアになっているトレーニング目標を使用して圧縮されたジェネレーター G をトレーニングできます。

学習目標は次のようにまとめられます。

これらの変更に基づいて、同じ圧縮フレームワークをさまざまなタイプの cGAN に適用できるようになりました。さらに、前述の疑似ペアを使用して学習すると、元のペアになっていないトレーニング設定と比較して、トレーニングがより安定し、より良い結果が得られます。

2. 教師の識別器から学ぶ

この研究ではジェネレータの圧縮に焦点を当てていますが、識別器 D は GAN にとって有用な情報を格納します。そのため、研究者たちは同じ識別器アーキテクチャを活用し、教師識別器からの事前トレーニング済みの重みを使用し、圧縮ジェネレータとともに識別器を微調整しました。

実験では、研究者らは、事前に訓練された識別器が学生ジェネレータの訓練を導くことができることを観察した。ランダムに初期化された識別子を使用すると、トレーニングが不安定になり、画像の品質が低下することがよくあります。この GAN の目標は次のように記述できます。

上記の式では、教師識別器 D′ からの重みを使用して、生徒識別器 D を初期化します。標準的なミニマックス オプティマイザーを使用して G と D をトレーニングします。

3. 中間特徴抽出

知識蒸留は、CNN モデルの圧縮で使用される一般的な方法です。出力層ロジットの分布を一致させることで、教師モデルからのダーク知識を生徒モデルに転送することができ、生徒モデルのパフォーマンスが向上します。ただし、条件付き GAN は通常、確率分布ではなく決定論的な画像を出力します。

上記の問題を解決するために、研究者は教師ジェネレータの中間表現を一致させました。中間層に含まれるチャネルの数が多いほど、提供できる情報量が多くなり、学生モデルが出力を超えて取得できる情報も増えます。蒸留の目標は次のとおりです。

ここで、G_t(x)とG′_t(x)は、生徒モデルと教師モデルのt番目に選択された層の中間特徴活性化であり、Tは層の数を表します。

4. 最適化目標を達成する

最終目標は次のように記述できます。

ハイパーパラメータ λ_recon と λ_distill は、各用語の重要性を制御します。

効率的な発電機設計空間

適切に設計された生徒アーキテクチャを選択することは、知識蒸留の最終的な効果にとって重要です。研究者は、教師モデルのチャネル数を単純に減らしても生徒モデルはコンパクトにならないことを発見しました。計算量の削減が 4 倍を超えると、パフォーマンスが大幅に低下します。

1. 畳み込み分解と層感度

既存のジェネレーターは通常、CNN 分類およびセグメンテーション設計に合わせて従来の畳み込みを採用しています。最近、いくつかの効率的な CNN 設計では、パフォーマンスと計算のバランスがより優れている畳み込みの分解形式 (深さ方向 + ポイント方向) が広く採用されています。研究者たちは、分解畳み込みは cGAN のジェネレータ設計にも使用できることを発見しました。

2. NASを使用してチャンネルを自動的にトリミングする

既存のジェネレーターは、すべてのレイヤーで手動で設計された（ほぼ均一な）数のチャネルを使用するため、冗長性が生じ、最適とは程遠い状態になります。圧縮効率をさらに向上させるために、研究者らはチャネルプルーニングを使用してジェネレーター内のチャネル幅を自動的に選択し、冗長性を減らして計算量を 2 倍に削減しました。このアプローチにより、チャネル数を細かく選択することができ、各畳み込み層に対して 8 の倍数から畳み込み層を選択できるため、MAC とハードウェアの並列処理のバランスをとることができます。

トレーニングとアーキテクチャ検索の分離

ワンショット NAS 手法に関する最近の研究に続いて、研究者らはモデル トレーニングとアーキテクチャ検索を切り離しました。まず、異なる数のチャネルを持つ「一度限りの」ネットワークがトレーニングされ、各サブネットワークは同じようにトレーニングされます。図 3 はフレームワーク全体を示しています。研究者は、元の教師ジェネレーターが1つのチャネルを持っていると仮定しています。与えられたチャネル数に対して

、「一度限りの」重みテンソルから最初のチャネルを抽出して、サブネットワークの重みネットワークを取得します。

各トレーニングステップでは、学習目標を使用して一定数のチャネルを持つサブネットワークがランダムにサンプリングされ、出力と勾配が計算され、抽出された重みが更新されます（式4）。抽出された最初のいくつかのチャネルはより頻繁に更新されるため、すべての重みの中でより重要な役割を果たします。

この「一度きりの」ネットワークがトレーニングされた後、研究者は検証セットで各サブネットワークのパフォーマンスを直接評価し、最適なサブネットワークを見つけました。 「一度限りの」ネットワークは重み共有によって徹底的にトレーニングされており、微調整は必要ありません。この結果は、モデルを最初からトレーニングした場合のパフォーマンスに近いものです。

この方法では、一度トレーニングするだけで、すべてのチャネルのすべての構成を評価し、検索結果に基づいて最適な構成を見つけるために、それ以上のトレーニングは必要ありません。もちろん、選択したアーキテクチャを微調整して、パフォーマンスをさらに向上させることも可能です。

実験結果

研究者らは、本論文のGAN圧縮フレームワークの一般化を検証するために、CycleGAN、Pix2Pix、GauGANという3つの条件付きGANモデルで実験を行った。使用される 4 つのデータセットは、Horse↔zebra、Edges→shoes、Cityscapes、および Map↔aerial photo です。

以下の表 1 は、上記 4 つのデータセットで CycleGAN、Pix2Pix、GauGAN モデルを圧縮した量子化結果を示しています。

表 1: Cityscapes データセットの mAP メトリック (高いほど良い) と他のデータセットの FID メトリックを使用した、3 つの条件付き GAN モデルの圧縮の定量評価。結果によると、この論文のGAN圧縮方法は、パフォーマンスがわずかに低下するものの、現在のSOTA条件付きGANをMACで7〜21倍、モデルサイズで5〜33倍圧縮できることがわかりました。 CycleGAN モデルの圧縮に関しては、この論文の GAN 圧縮方法は、以前の CycleGAN 固有の共進化方法よりもはるかに優れています。

パフォーマンスと計算のトレードオフ

この方法は、より高い圧縮率を実現するだけでなく、さまざまなモデル サイズのパフォーマンスも向上させます。下の図 6 は、pix2pix モデルにおけるさまざまなデータセットのパフォーマンスと計算のトレードオフを示しています。

図 6: Cityscapes および Edges→Shoes データセットにおける pix2pix のトレードオフ曲線。プルーニングと蒸留の方法は、大規模なモデルを最初からトレーニングするよりも優れていますが、モデルが大幅に圧縮されている場合はパフォーマンスが低下します。

エフェクト表示

下の図4は、この方法を使用して得られた効果を示しています。図は、それぞれ入力データ、標準出力、元のモデル出力、圧縮モデル出力を示しています。図からわかるように、研究者が提案した方法では、圧縮率が高くなっても、出力画像の視覚的な信頼性を維持できます。

図 4: Cityscapes、Edges→Shoes、Horse→Zebra データセットの結果の比較。

ハードウェア推論アクセラレーション

現実世界のインタラクティブ アプリケーションでは、コンピューティング消費を削減することよりも、ハードウェア デバイスでの推論を高速化することの方がはるかに重要です。以下の表 2 に示すように、提案手法の実際のアプリケーションでの有効性を検証するために、研究者はさまざまなコンピューティング性能を持つデバイス上で圧縮モデルの推論速度をテストしました。

表 2: NVIDIA Jetson AGX Xavier、NVIDIA Jetson Nano、1080Ti GPU、Xeon CPU で測定されたメモリとレイテンシの削減。

結論は

この記事では、Han Song 氏のチームが提案した一般的な圧縮フレームワークにより、条件付き GAN におけるジェネレーターの計算コストとモデル サイズが大幅に削減され、知識の蒸留と自然なアーキテクチャの検索を通じてトレーニングの安定性とモデル効率が向上します。実験では、本論文で提案された GAN 圧縮方法は、視覚的な品質を維持しながら、いくつかの条件付き GAN モデルを圧縮できることが示されています。今後の研究は、モデルの遅延を減らし、ビデオモデルを生成するための効率的なフレームワークを構築することに重点を置く予定だと彼らは述べた。