gemmlowpライブラリのソースコードをデバッガを使って理解する

今回もAIの量子化について学んでいきます。GoogleのQAT量子化の論文に出てくる gemmlowp ライブラリについて見ていきます。

前回は、gemmlowp のサンプルコード（doc/quantization_example.cc）を実行しました。長い実行ログが出力されて、その内容は量子化計算のチュートリアルのようになっていたので、その説明をしました。

今回は、実行結果のログだけでは十分に理解できなかった部分を、デバッガを使って、実行結果とソースコードを見ながら、理解を深めていきます。

それでは、やっていきます！

はじめに
行列の積を量子化で計算する
行列積を量子化で計算した結果を逆量子化する
量子化で行列積を計算した結果と浮動小数点演算で計算した結果を比較する
おわりに

はじめに

「AIモデルの量子化」の記事一覧です。良かったら参考にしてください。

AIモデルの量子化の記事一覧

・量子化：論文 Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference を読む（１）
・量子化：論文 Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference を読む（２）
・TensorFlow Lite Python で量子化モデルをRaspberry Pi 4で動かす
・TensorFlow Lite C++ で量子化モデルをRaspberry Pi 4で動かす
・Raspberry Pi 4 の TensorFlow Lite C++ を VSCode でリモートデバッグする
・gemmlowpライブラリのソースコードを読んで量子化を学ぶ（サンプルソースの実行）
・gemmlowpライブラリのソースコードをデバッガを使って理解する ← 今回

gemmlowp のリポジトリは、以下です。

github.com

それではやっていきます！

行列の積を量子化で計算する

前回のおさらいです。

サンプルソースは、スケールとゼロポイントを求めて、LHS と RHS を量子化していました。その後、量子化した LHS と RHS の積を計算していました。

今回は、その行列積の計算が少し難しいので、デバッガを使って順番に見ていきます。

行列の積を量子化で計算するソースコード

まず、該当のソースコードを示します。パッと見て、全くわかりません（笑）。

なので、少しずつ見ていきます。

  gemmlowp::OutputStageQuantizeDownInt32ByFixedPoint quantize_down_stage;
  quantize_down_stage.result_offset_after_shift = result_offset;
  quantize_down_stage.result_fixedpoint_multiplier = quantized_multiplier;
  quantize_down_stage.result_shift = right_shift;
  gemmlowp::OutputStageSaturatingCastToUint8 saturating_cast_stage;
  const auto& output_pipeline =
      std::make_tuple(quantize_down_stage, saturating_cast_stage);

  auto actual_uint8_result_map = actual_uint8_result.Map();
  gemmlowp::GemmContext gemm_context;
  gemmlowp::GemmWithOutputPipeline<std::uint8_t, std::uint8_t,
                                   gemmlowp::DefaultL8R8BitDepthParams>(
      &gemm_context, uint8_lhs.ConstMap(), uint8_rhs.ConstMap(),
      &actual_uint8_result_map, lhs_offset, rhs_offset, output_pipeline);

  std::cout << "Quantized uint8 result matrix obtained by quantized "
            << "multiplication:\n"
            << actual_uint8_result << std::endl;

まず、最初は gemmlowp::OutputStageQuantizeDownInt32ByFixedPoint という構造体の変数を定義しています。

構造体の定義は以下です。

struct OutputStageQuantizeDownInt32ByFixedPoint {
  std::int32_t result_fixedpoint_multiplier;
  std::int32_t result_shift;
  std::int32_t result_offset_after_shift;
};

次に、構造体のメンバに値を設定していますが、その値を計算している部分を示します。

real_multiplier は、LHS、RHS、計算結果のスケールから計算しています。論文では、 $M := \dfrac{S_1S_2}{S_3}$ と書かれていました。

real_multiplier を使って、quantized_multiplier と right_shift を計算しているようです。

  const float real_multiplier = lhs_qparams.scale * rhs_qparams.scale / result_qparams.scale;
  std::int32_t quantized_multiplier;
  int right_shift;
  QuantizeMultiplierSmallerThanOne(real_multiplier, &quantized_multiplier, &right_shift);

では、QuantizeMultiplierSmallerThanOne() を見ていきます。

QuantizeMultiplierSmallerThanOne()

まず、ソースコードを示します。

先頭の assert() は、論文に $M$ は $(0, 1)$ と書かれていた部分ですね。

まず、real_multiplier を $[0.5, 1)$ の範囲になるまで、s 回だけ2倍（左シフト）します。

次は、1 << 31 とかけて、round（四捨五入）しています。これは、小数部が 31bit の固定小数点数に変換しています。残りの 1bit は整数部（もしくは、符号部）です。

求まった q が quantized_multiplier で、s が right_shift になります。それ以外にもいろいろしてますが、チェックだったり、丸め誤差などを考慮した調整などです。

void QuantizeMultiplierSmallerThanOne(float real_multiplier,
                                      std::int32_t* quantized_multiplier,
                                      int* right_shift) {
  assert(real_multiplier > 0.f);
  assert(real_multiplier < 1.f);
  int s = 0;

  while (real_multiplier < 0.5f) {
    real_multiplier *= 2.0f;
    s++;
  }

  std::int64_t q = static_cast<std::int64_t>(std::round(real_multiplier * (1ll << 31)));
  assert(q <= (1ll << 31));

  if (q == (1ll << 31)) {
    q /= 2;
    s--;
  }
  assert(s >= 0);
  assert(q <= std::numeric_limits<std::int32_t>::max());
  *quantized_multiplier = static_cast<std::int32_t>(q);
  *right_shift = s;
}

実際に値をデバッガで見ると、real_multiplier は 0.00436593033、quantized_multiplier は 1200097792、right_shift は 7 でした。

これは、real_multiplier に $2^{38}$ をかけたことと同じです。

これで、構造体 quantize_down_stage のメンバの値が計算できました。

gemmlowp::GemmWithOutputPipeline()

あとは gemmlowp::GemmWithOutputPipeline() で量子化で行列積を計算するだけですが、実際にデバッガで読み進めていったところ、非常に複雑で、それを全て説明するのは難しいと思いました。

そこで、計算の流れを説明します。

量子化した LHS と RHS は以下です。

Quantized uint8 LHS matrix:
208     236     0       238
3       214     255     29

Quantized uint8 RHS matrix:
152     51      244
60      26      255
0       127     246
127     254     247

この計算結果が、以下です。

Quantized uint8 result matrix obtained by quantized multiplication:
168     115     255
0       66      151

計算結果の 168 は、行列積なので、208 * 152 + 236 * 60 + 0 * 0 + 238 * 127 のように計算されます。計算を続けると、31616 + 14160 + 0 + 30226 = 76002 となります。

これは、論文の 2.3 章の式 (7) の $\sum\limits_{j=1}^{N}q_1^{(i, j)}q_2^{(j, k)}$ になります。

$q_3^{(i, k)} = Z_3 + M \left( NZ_1Z_2 - Z_1a_2^{(k)} - Z_2a_1^{(i)} + \sum\limits_{j=1}^{N}q_1^{(i, j)}q_2^{(j, k)} \right) \tag7$

$a_2^{(k)} := \sum\limits_{j=1}^{N}q_2^{(j, k)}, \hspace{0.5em} a_1^{(i)} := \sum\limits_{j=1}^{N}q_1^{(i, j)} \tag8$

式 (7) を計算すれば、168 になるはずです。やってみます。

$N$ は 4 で、 $Z_1$ は 113、 $Z_2$ は 114、 $Z_3$ は 118、 $a_2^{(k)}$ は 152 + 60 + 0 + 127 = 339、 $a_1^{(i)}$ は 208 + 236 + 0 + 238 = 682 です。

式 (7) の括弧内は、4 * 113 * 114 - 113 * 339 - 114 * 682 + 76002 = 51528 - 38307 - 77748 + 76002 = 11475 となります。

あとは、 $M$ をかけて、 $Z_3$ を足してやればいいわけですが、 $M$ は小数であり、整数だけで計算したいので、代わりに、以下を使います。

$M = 2^{-n}M_0 \tag6$

quantized_multiplier は 1200097792、right_shift は 7 でした。

11475 * 1200097792 = 13,771,122,163,200 で、 $2^{31}=2147483648$ で割って、 $2^7=128$ で割ると $50$ になります。

あとは、 $Z_3$ の $118$ を足して、 $168$ となります。

なぜ、 $2^{31}$ で割るかというと、 $M_0$ は 31bit の固定小数点数だからであり、元に戻す必要があるからです。

また、 $2^7$ で割る理由は、real_multiplier を $[0.5, 1)$ の範囲にするために right_shift の数だけ左シフトしていたので、元に戻す必要があるからです。

もっと簡単に言うと、 $M$ に $2^{38}$ をかけて、括弧内の結果（11475）と掛け算して、その結果を $2^{38}$ で割っただけです。

なぜ、こんなややこしいことをするのかというと、浮動小数点演算の結果となるべく近づけたいからです。つまり、real_multiplier の小数点以下の値をなるべく演算に反映したいからです。

行列積を量子化で計算した結果を逆量子化する

前回の宿題で、手計算で逆量子化すると、サンプルソースの結果と異なっていた件がありました。

サンプルソースが逆量子化した結果は以下で、手計算すると、r = 0.0107 * (168 - 118) = 0.535 でした。

Here is the actual float product (LHS * RHS) matrix obtained by dequantizing the above uint8 result, i.e. as far as we are concerned, the ACTUAL RESULT:
0.533   -0.032  1.46
-1.26   -0.554  0.352

デバッガで見ると、スケールは 0.0106628919 だったので、それを使うと、r = 0.0107 * (168 - 118) = 0.533144595 となります。サンプルソースの結果と同じになりました。

量子化で行列積を計算した結果と浮動小数点演算で計算した結果を比較する

前回の宿題で、手計算で差分を計算したら、サンプルソースの結果と異なっていた件がありました。

サンプルソースの差分は以下で、手計算すると、0.533 - 0.534 = -0.001 でした。

Difference between ACTUAL and REFERENCE float results:
-0.000675       0.00764 -0.000674
-0.000674       0.0022  0.00369

デバッガで見ると、浮動小数点演算した結果は、0.533819675 で、量子化で計算した結果は 0.533144595 なので、引くと、-0.00067508 となり、結果が一致しました。

おわりに

今回は、gemmlowp ライブラリのサンプルソースの実行を細かく見ていきました。TensorFlow Lite C++ が、どのように量子化して積和演算を行っているかが、よく分かりました。

最後になりましたが、エンジニアグループのランキングに参加中です。

ランキング参加中

エンジニアグループ

気楽にポチッとよろしくお願いいたします🙇

今回は以上です。

最後までお読みいただき、ありがとうございました。

daisukeの技術ブログ

AI、機械学習、最適化、Pythonなどについて、技術調査、技術書の理解した内容、ソフトウェア/ツール作成について書いていきます