TensorRT の EntropyCalibrator の観察

ディープニューラルネットワーク (DNN) の量子化は、 主に、モデルの Weight と Activation を単精度浮動小数点 (FP32) から 8ビット整数 (INT8) へ 変換することを指す。 この変換は、浮動小数点のスケール $s$ と整数のゼロ点 $z$ を用いて表される。

本稿では、このスケール $s$ の決定のためのキャリブレーションの手法の一つ、 Entropy についてそのアルゴリズムと振る舞いについて書く。

なお、量子化そのものについては、文献 [1] で Quantization Aware Training (QAT) や 8 ビット以下の量子化などを含めた俯瞰的なサーベイが行われている。 また、文献 [2] では、各キャリブレーション手法の評価を含む、実践的な内容が述べられている。

Calibration のアルゴリズム

代表的な Calibration のアルゴリズムとして、以下の 3 種類が挙げられる。[2]

なかでも、TensorRT においてデフォルトの手法となっている Entropy についてその振る舞いを理解したい。

Entropy

GTC2017 での NVIDIA による発表資料 [3] をもとに、どのように閾値（スケール）が決定されるかを観察する。

FP32 から INT8 への変換を、再エンコーディングと捉え、その間の情報の損失を最小化することを考える。 2 つのエンコーディングを確率分布と見れば、その2つの確率分布の距離は Kullback-Leibler divergence を用いて表せる。 KL-divergence により、与えられたエンコーディングを近似したときの、情報損失の量を測ることができる。 したがって、情報損失が最小となるような近似を作るのが目的となる。

アルゴリズム

  1. まず、参照分布 $P$ を作る。
    1. $H$ の $i-1$ 番目までのビンを使い、 $$P = [m_0, \ldots, m_{i-1}]$$ とする。
    2. $H$ の $i$ 以降のビンの合計を計算し、これを外れ値 $O$ とする。 $$O = \sum_{k=i}^{N-1} m_k$$
    3. 外れ値 $O$ を $P'$ の末尾に足したものを参照分布 $P$ とする。 $$ P(i-1) \leftarrow P(i-1) + O $$
  2. 次に、候補分布 $Q$ を作る。
    1. $[m_0, \ldots, m_{i-1}]$ を $2^8$ レベルに量子化する。 つまり、 $i$ 個のビンに分割されたヒストグラムを $2^8$ 個のビンに分割し直す。
    2. 量子化したヒストグラムを再度 $i$ 個に展開し、 候補分布 $Q$ とする。 このとき、もとのヒストグラム $H$ で $0$ 個であったビンについては、$0$ を保持する。
  3. $P, Q$ それぞれをその合計で割り、確率分布とする。 $$ P \leftarrow P / \sum{P} $$ $$ Q \leftarrow Q / \sum{Q} $$
  4. 得られた参照分布 $P$ と候補分布 $Q$ について、KL-divergence を計算する。 $$ KL(P||Q) = \sum_{k=0}^{i-1} P(k) \log\frac{P(k)}{Q(k)} $$

以上から、 $KL(P||Q)$ が最小となる $i$ を求める。

TensorRT の EntropyCalibrator での振る舞いの観察

公式の実装を参考に、 実際にどのように上記のループが行われるのか観察してみる。

実験設定

結果

下図に、$i = 2^2, \ldots, 2^4$ について、参照分布 $P$（青）と候補分布 $Q$ （橙）をプロットした。

まとめ

本稿では、 TensorRT での EntropyCalibrator の中身の挙動について アルゴリズムを振り返り、簡単な例での中身の挙動を観察した。

Reference

#Deep Learning #Quantization #TensorRT #EntropyCalibrator