本文首發於公眾號「AI小男孩」，歡迎大伙過來砸場！
在之前的文章中提到過可以把 ReLU 合並到 Conv 中加速量化推理，當時只是用一個例子簡單介紹一下過程，邏輯上存在一些漏洞。本文打算從數學上深入剖析一下其中的原理，並進一步擴展到其他激活函數，看看在網絡量化中激活函數一般是怎么處理的。
溫故知新
為了簡單起見，假設我們是量化到 uint8 的數值范圍「即0~255」。回憶一下量化的基本公式「我在之前的文章中多次強調這幾個公式，它們非常重要」
再簡單重復一下符號的含義，\(r\) 表示實數，\(q\) 表示量化后的定點數，\(S\) 和 \(Z\) 分別是是 scale 和 zero point。
注意，這次我對 \(q\) 單獨加了一個 clip 操作，在之前的文章中，這一步在公式中被我省略了，不過在實際量化的時候，這一步是必須的，否則會有數值溢出的危險。
現在，假設激活函數為 \(f(x)\)，應用到實數域上是這個樣子：
那么把 (1) 式代入后可以得到量化的公式：
這就是量化時處理所有激活函數的總口訣，別看它平平無奇，但話越少，信息量越多。下面，我們就看看針對具體的激活函數，怎么運用這個公式。
ReLU
ReLU 是一個非常簡單的函數，它除了把小於 0 的數值截斷外，甚至不做任何操作：
如果把上面的函數 \(f\) 替換成 ReLU 的公式，就可以得到：
把 (1) 式代入就變成：
換算一下可以得到：
這是量化 ReLU 最通用的運算，其中 \(\frac{S_1}{S_2}\) 可以通過之前文章講的定點數 + bitshift 來實現。
需要重點指出的是，ReLU 之后，\(Z_2\) 永遠等於 0。因為 ReLU 會把實數域上小於 0 的數全部截斷為 0，此時去統計實數域的范圍，可以發現是 0~a，而我們量化的數值范圍是 0～255，為了保證零點對齊，因此 \(Z_2\) 只能取 0。
當然啦，具體實現上沒有必要完全按照 (8) 式來操作。一來公式內的 scale 操作過於麻煩還掉精度，二來 ReLU 本身是有明確的物理意義的，那就是把小於零點的數值截斷，其余不變。這個意義在量化里面依然成立。
因此，我們其實可以用一種更簡潔明了的方式來實現量化的 ReLU：
如果是使用這個公式，那 ReLU 前后的 scale 和 zeropoint 是要保持一致的，這一點可以從 ReLU 本身的物理含義出發得出。
tflite 里面就是用了這個簡化的公式來實現 ReLU 的功能「下面這段代碼參考自https://github.com/tensorflow/tensorflow/blob/r1.15/tensorflow/lite/kernels/internal/reference/reference_ops.h#L214」：
template <typename T>
inline void ReluX(const tflite::ActivationParams& params,
const RuntimeShape& input_shape, const T* input_data,
const RuntimeShape& output_shape, T* output_data) {
gemmlowp::ScopedProfilingLabel label("Quantized ReluX (not fused)");
const int flat_size = MatchingFlatSize(input_shape, output_shape);
const T max_value = params.quantized_activation_max;
const T min_value = params.quantized_activation_min;
for (int i = 0; i < flat_size; ++i) {
const T val = input_data[i];
const T clamped =
val > max_value ? max_value : val < min_value ? min_value : val;
output_data[i] = clamped;
}
}
可以看出，這個量化的 ReLU 和浮點數版本的 ReLU 邏輯上幾乎沒有區別。
ReLU如何勾搭上Conv
其實不止是 Conv，全連接層 FC 等也可以和 ReLU 合並。我們來看看為什么。
同樣地，假設一個卷積操作為 \(r_3=\sum_{i}^N r_1^i r_2^i\)，按照之前文章的描述，量化后的公式為：
現在，\(q_3\) 進入 ReLU 進行運算得到 \(q_4\)，按照上面的推算可以得出：
換算一下得到：
到這里，這個式子仍然是 ReLU 的形式。換句話說，我們仍然要走兩個分支來計算函數的結果。
但是，如果要把 ReLU 合並到 Conv 中，就必須得用 Conv 的運算來代替這個分支運算。換句話說，\(q_4\) 無論跑哪個分支，都必須可以用 \(\frac{S_1S_2}{S_4}\sum_{i}^N (q_1-Z_1)(q_2-Z_2)+Z_4\) 直接計算出來，我們才能實現 Conv 和 ReLU 的合並。
這時，就要用到量化中的 clip 操作了。上面式子 (12)，其實更嚴格的寫法應該是：
前面說了，\(Z_4=0\)。如果 \(q_3 < Z_3\)，那么等價地 \(\sum_{i}^N (q_1-Z_1)(q_2-Z_2)<0\)，此時會跑第二個分支得到 \(q_4=Z_4\)。但是，由於有 clip 操作，在這種情況下，\(q_4=clip(\frac{S_1S_2}{S_4}\sum_{i}^N (q_1-Z_1)(q_2-Z_2)+Z_4, 0, 255)=0=Z_4\)，因此，我們發現，無論跑哪個分支，最后都可以統一用下面這個式子來表示：
而這個公式的意義相當於：我們計算出 ReLU 之后的 \(S\) 和 \(Z\)，然后把這個 \(S\) 和 \(Z\) 對應到 Conv 的輸出，這樣一來，ReLU 的運算就合並到 Conv 里面了。
正如我前面提到的，ReLU 除了做數值上的截斷外，其實沒有其他操作了，而量化本身自帶截斷操作，因此才能把 ReLU 合並到 Conv 或者 FC 等操作里面。
LeakyReLU
有讀者可能覺得，ReLU 本身的操作很簡單，為什么還得用 (8) 式這種繞彎路的方式說一大堆。那是因為 ReLU 本身的性質可以讓我們抄近道，如果換成其他函數，這個流程就繞不過去了。
不信來看看 LeakyReLU 是怎么量化的。
LeakyReLU 的公式可以表示成：
這里面的 \(\alpha\) 是我們事先指定的數值，一般是 0~1 之間的小數。
同樣地，我們按照文章最開始的總口訣，即公式 (3)(4)，來逐步分析這個函數。把原來的函數 \(f\) 替換成 LeakyReLU，可以得到：
把 (1) 式代入：
換算一下得到：
此時，由於有 \(\alpha\) 的存在，這兩個分支就無法像 ReLU 一樣進行合並，自然也就無法整合到 Conv 等操作內部了。
在 tflite 中是將 \(\alpha\) 轉換為一個定點數再計算的。具體地，假設 \(\alpha_q=clip(round(\frac{\alpha}{S_1}+Z_1), 0, 255)\)，可以得到：
具體代碼如下「參考自https://github.com/tensorflow/tensorflow/blob/r1.15/tensorflow/lite/kernels/activations.cc#L248」：
TfLiteStatus LeakyReluPrepare(TfLiteContext* context, TfLiteNode* node) {
TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
const TfLiteTensor* input = GetInput(context, node, 0);
TfLiteTensor* output = GetOutput(context, node, 0);
TF_LITE_ENSURE_EQ(context, input->type, output->type);
LeakyReluOpData* data = reinterpret_cast<LeakyReluOpData*>(node->user_data);
if (output->type == kTfLiteUInt8) {
const auto* params =
reinterpret_cast<TfLiteLeakyReluParams*>(node->builtin_data);
// Quantize the alpha with same zero-point and scale as of input
data->q_alpha = static_cast<uint8_t>(std::max<float>(
std::numeric_limits<uint8_t>::min(),
std::min<float>(std::numeric_limits<uint8_t>::max(),
std::round(input->params.zero_point +
(params->alpha / input->params.scale)))));
double real_multiplier =
input->params.scale * input->params.scale / output->params.scale;
QuantizeMultiplierSmallerThanOneExp(
real_multiplier, &data->output_multiplier, &data->output_shift);
}
return context->ResizeTensor(context, output,
TfLiteIntArrayCopy(input->dims));
}
這段代碼主要是做一些准備工作，把 \(\alpha_q\) 和 \(\frac{S_1S_1}{S_2}\) 等變量事先計算好。
函數本身的具體操作如下「參考自https://github.com/tensorflow/tensorflow/blob/r1.15/tensorflow/lite/kernels/internal/reference/reference_ops.h#L242」：
template <typename T>
inline void QuantizeLeakyRelu(const LeakyReluParams& params, T q_alpha,
const RuntimeShape& input_shape,
const T* input_data,
const RuntimeShape& output_shape,
T* output_data) {
gemmlowp::ScopedProfilingLabel label("LeakyRelu (not fused)");
const int flat_size = MatchingFlatSize(input_shape, output_shape);
static const int32 quantized_min = std::numeric_limits<T>::min();
static const int32 quantized_max = std::numeric_limits<T>::max();
static const int32 alpha_value = q_alpha - params.alpha_offset;
for (int i = 0; i < flat_size; ++i) {
const int32 input_value = input_data[i] - params.input_offset;
if (input_value >= 0) {
output_data[i] = input_data[i];
} else {
const int32 unclamped_output =
params.output_offset + MultiplyByQuantizedMultiplierSmallerThanOneExp(
input_value * alpha_value,
params.output_multiplier,
params.output_shift);
const T clamped_output =
std::min(quantized_max, std::max(quantized_min, unclamped_output));
output_data[i] = static_cast<uint8>(clamped_output);
}
}
}
代碼里面的 input_value 就是公式 (19) 里面的 \(q_1-Z_1\)，tflite 會根據 input_val 的數值情況分兩個分支運行，這個過程和 (19) 基本一致。
眼尖的讀者可能發現，為啥 \(q_1>Z_1\) 這個分支，代碼里面好像直接令 \(q_2=q_1\) 了，這跟公式 (19) 描述的好像不一樣啊。哈哈，這個地方我也暫時不明白，了解詳情的讀者請教教我，或者我之后弄懂再補充一下。
非線性函數
對於類 ReLU 函數來說，其實還都是分段線性的，那遇到非線性的函數「比如 sigmoid、tanh」又該怎么量化呢？從 gemmlowp 的文檔來看，這些函數其實是用定點運算來近似浮點的效果。這部分內容觸及到我的知識盲區，所以就不給大家做深入介紹了，感興趣的讀者可以看一下 gemmlowp 的源碼進一步了解：https://github.com/google/gemmlowp/blob/master/fixedpoint/fixedpoint.h。
雖然我對里面的原理了解不多，但還是有一點點落地的經驗。我曾經用高通驍龍的 SNPE 工具量化了 tanh 函數，但在 DSP 上跑定點運算的時候，發現耗時比在 GPU 上跑浮點運算滿了 100 倍左右。
因此對於有落地需求的同學來說，我的建議是網絡中盡量不要包含這類非線性函數。如果實在要有的話，要么嘗試把網絡拆成幾塊，一些跑定點，一些跑浮點，要么就是用一些線性函數來近似這些非線性函數的效果。
總結
這篇文章主要講了網絡量化中如何處理激活函數，並從數學上進一步剖析為何 ReLU 可以和 Conv 等操作合並。
你可能已經發現，網絡量化這個課題跟底層的實現聯系非常緊密，比如涉及到 gemmlowp、neon 等底層函數庫等。有讀者會說：我只想老老實實研究算法，對這些底層的運算不了解也沒興趣了解啊！
對於這部分讀者，其實也不用焦慮，誠然，網絡量化對底層的聯系相比其他深度學習算法而言更加緊密，但對於頂層的算法開發人員，只需要大概知道底層是怎么運行的就可以，而把更多的精力放在對量化算法的改進上。當然啦，如果想成為一流的網絡量化專家，熟悉底層還是很有必要的，否則你怎么知道未來算法的發展趨勢呢？
PS. 最近陸續填了一些坑，之后應該會介紹一些更前沿且對落地比較友好的論文和技術了。感謝大家在我斷更這么久后依然不離不棄。