好久沒更新了，一方面是因為工作繁忙，另一方面主要是懶。
之前寫過幾篇關於神經網絡量化的文章，主要是對 Google 量化論文以及白皮書的解讀，但有一些細節的問題當時沒有提及。這篇文章想補充其中一個問題：關於 ElementwiseAdd (簡稱 EltwiseAdd) 和 Concat 的量化。
EltwiseAdd量化
EltwiseAdd 的量化主要是在論文的附錄里面提及的。過程不是太復雜，如果了解量化的基本原理的話，完全可以自己推導出來。
回憶一下量化的基本公式：
(看不懂的可以再參考一下我之前的文章)
這里面 \(r\) 是實數域中的數值 (一般是 float)，\(q\) 則是量化后的整型數值 (常用的是 int8)。
EltwiseAdd 就是對兩個 tensor 的數值逐個相加。假設兩個 tensor 中的數值分別是 \(r_1\)、\(r_2\)，相加得到的和用 \(r_3\) 表示，那全精度下的 EltwiseAdd 可以表示為：
用量化的公式代入進去后可以得到：
稍作整理可以得到：
注意，這里有兩個 scale 運算需要轉換為定點小數加一個 bitshift 的運算 (具體做法見之前的文章)。除了需要對輸出按照 \(\frac{S_1}{S_3}\) 放縮外，其中一個輸入也需要按照 \(\frac{S_2}{S_1}\) 進行放縮，這一步就是論文中提到的 rescale。
這一部分的代碼我就不准備在 pytorch 中實現了，畢竟這個模塊的量化最主要的就是統計輸入跟輸出的 minmax，因此訓練代碼幾乎沒什么內容，主要的工作都是在推理引擎實現的。因此這篇文章我會摘取 tflite 中部分實現簡單說明一下。
下面是 tf1.5 中我摘取的部分關於 EltwiseAdd 的量化實現，對應的鏈接是https://github.com/tensorflow/tensorflow/blob/r1.15/tensorflow/lite/kernels/internal/reference/add.h#L53：
inline void AddElementwise(int size, const ArithmeticParams& params,
const uint8* input1_data, const uint8* input2_data,
uint8* output_data) {
// ......此處省略若干無關代碼
for (int i = 0; i < size; ++i) {
const int32 input1_val = params.input1_offset + input1_data[i];
const int32 input2_val = params.input2_offset + input2_data[i];
const int32 shifted_input1_val = input1_val * (1 << params.left_shift);
const int32 shifted_input2_val = input2_val * (1 << params.left_shift);
const int32 scaled_input1_val =
MultiplyByQuantizedMultiplierSmallerThanOneExp(
shifted_input1_val, params.input1_multiplier, params.input1_shift);
const int32 scaled_input2_val =
MultiplyByQuantizedMultiplierSmallerThanOneExp(
shifted_input2_val, params.input2_multiplier, params.input2_shift);
const int32 raw_sum = scaled_input1_val + scaled_input2_val;
const int32 raw_output =
MultiplyByQuantizedMultiplierSmallerThanOneExp(
raw_sum, params.output_multiplier, params.output_shift) +
params.output_offset;
const int32 clamped_output =
std::min(params.quantized_activation_max,
std::max(params.quantized_activation_min, raw_output));
output_data[i] = static_cast<uint8>(clamped_output);
}
}
這里面有個函數 MultiplyByQuantizedMultiplierSmallerThanOneExp，它的主要作用是調用 gemmlowp 中的函數將乘以 scale 的浮點運算轉換為乘以一個定點小數加 bitshift 的操作，由於涉及比較多底層操作，不在本文討論之內。
整段代碼的邏輯和上文分析的基本類似，首先是對輸入加 offset 操作，對應公式中的 \(q_i-Z_i\)，然后分別對兩個輸入乘以 scale，那按照上文的描述，一般來說只有一個輸入需要進行 rescale 操作，另一個輸入的 scale 其實是 1。在對兩個輸入相加后得到輸出 (代碼中的 raw_sum)，會按照同樣的方式對輸出進行 scale 放縮並加上 offset，最后再 clamp 到 uint8 的數值范圍內。
Concat量化
Concat 可以采用和 EltwiseAdd 類似的操作，對其中一個輸入進行 rescale 后再 concat，最后再對輸出進行 rescale，參考如下推導：
代入量化公式：
整理后得到：
不過 rescale 本身是存在精度損失的，而 Concat 嚴格來說是一個無損的操作 (concat 其實就是內存拷貝而已)，因此論文建議統一輸入輸出的 scale 來避免 rescale：
不過我始終想不通要如何在沒有 rescale 的情況下統一輸入輸出的 scale。論文中也沒有提及相關的實現，很多細節只能到 tflite 的源碼中查找。
可以明確的一點是，output 的 minmax 可以通過取兩個輸入的最小 min 和最大 max 來確定。那無非存在兩種情況：1. 其中一個輸入的 minmax 覆蓋了整個范圍，即輸出的 minmax 完全由某一個輸入確定；2. minmax 分別來自兩個輸入，即一個輸入的 min 和 另一個輸入的 max 確定輸出的 minmax。
為了了解 Google 到底怎么處理 concat 量化，我稍微翻了下 tf1.5 中對於量化 concat 的實現。
下面是在源碼中找到的部分代碼注釋：
// There are two inputs for concat, "input0" and "input1". "input0" has [0, 5]
// as min/max and "input1" has [0, 10] as min/max. The output "output" for
// concat has [0, 10] as min/max.
// After applyging QuantizeModel(), "input0" will have a requant op added, along
// with a tensor "input0_reqaunt" that has [0, 10] as min/max. So the topology
// becomes:
// input0 -> requant -> input0_requant \
// concat - output
// input1 /
這段注釋說的是上面的情況 1，即其中一個輸入的 minmax 覆蓋了整個范圍。這種情況下，tflite 的做法是將 range 較小的輸入進行 requant，即根據大 range 的 minmax，來重新量化這個輸入。
那具體怎么 requant 呢？這里需要在另一段代碼中找細節：
inline void ConcatenationWithScaling(const ConcatenationParams& params,
const RuntimeShape* const* input_shapes,
const uint8* const* input_data,
const RuntimeShape& output_shape,
uint8* output_data) {
....
const float inverse_output_scale = 1.f / output_scale;
uint8* output_ptr = output_data;
for (int k = 0; k < outer_size; k++) {
for (int i = 0; i < inputs_count; ++i) {
const int copy_size = input_shapes[i]->Dims(axis) * base_inner_size;
const uint8* input_ptr = input_data[i] + k * copy_size;
if (input_zeropoint[i] == output_zeropoint &&
input_scale[i] == output_scale) {
memcpy(output_ptr, input_ptr, copy_size);
} else {
const float scale = input_scale[i] * inverse_output_scale;
const float bias = -input_zeropoint[i] * scale;
for (int j = 0; j < copy_size; ++j) {
const int32_t value =
static_cast<int32_t>(std::round(input_ptr[j] * scale + bias)) +
output_zeropoint;
output_ptr[j] =
static_cast<uint8_t>(std::max(std::min(255, value), 0));
}
}
output_ptr += copy_size;
}
}
}
這里只貼了其中比較關鍵的實現，鏈接：https://github.com/tensorflow/tensorflow/blob/r1.15/tensorflow/lite/kernels/internal/reference/reference_ops.h#L1164。
具體做法是這樣的：如果輸入的 scale、zeropoint 和輸出不一樣，那么就對該輸入按照輸出的 scale 和 zeropoint 重新 requant，表示成公式的話是這樣子的：
對比上面公式 (7)，我發現這他喵不就是對輸入 \(q_1\) 進行 rescale 嗎？而且，上面這段代碼不會區分 \(q_1\)、\(q_2\)，只要發現輸入的 scale 和 zeropoint 和輸出對不上，就會對任何一個輸入進行 requant。
換言之，量化 concat 可以用公式表示為：
總結
這篇文章是對網絡量化中 EltwiseAdd 和 Concat 兩個操作的補充，由於有 rescale 以及 requant 的存在，這兩個運算相比 float 而言，計算量反而更大，而且可能導致精度上的損失。因此在量化網絡的時候，需要關注這兩個函數的輸入 range 不要相差太大，以避免精度損失過大。
參考
PS: 之后的文章更多的會發布在公眾號上，歡迎有興趣的讀者關注我的個人公眾號：AI小男孩，掃描下方的二維碼即可關注
