深度剖析多任務模型 QAT 策略

本文為筆者個人見解，如有不同意見歡迎評論

為了節省端側計算資源以及簡化部署工作，目前智駕方案中多采用動靜態任務融合網絡，地平線也釋放了 Lidar-Camera 融合多任務 BEVFusion 參考算法。這種多任務融合網絡的浮點訓練策略可以簡述為：

首先在大量數據的條件下完成多任務模型 backbone 部分的浮點訓練；

然后，固定 backbone 的權重，分別接多個 task head 進行單獨的訓練。

在這種浮點訓練策略下，QAT（量化感知訓練）里的 calibration（校準）和量化訓練策略跟常規的單 task 模型差別較大。常規的單 task 模型一般就是用那種比較固定、普遍適用的訓練辦法，不過碰到復雜多變又有特定需求的情況，它的局限性就慢慢顯現出來了。

這篇文章會結合具體的場景，對 calibration 和量化訓練策略進行分析，然后提出一些筆者個人獨特的看法，希望相關領域的研究和實踐提供點有用的參考和啟發。

本文將以具有兩個任務頭 task_head1 和 task_head_2 的多任務模型為例進行描述。

此步驟的前提是模型已完成**浮點訓練。**

在進行后續操作之前，必須確保模型已經成功地完成了浮點訓練。只有在這個前提條件得到滿足的情況下，才能保證后續的工作能夠順利進行，并且得到準確和可靠的結果。

首先對骨干網絡（backbone）進行校準/量化感知訓練（calibration/qat），在滿足量化精度要求后，保存校準/量化感知訓練的權重（calib/qat 權重）。

驗證 step1 中 backbone 在部署 head 上的精度，具體操作是對 backbone 進行偽量化處理，而 task_head1 和 task_head2 保持浮點計算，然后在驗證集上測試這兩個 head 的精度。

若 step2 驗證出的 task_head1 精度不符合預期，則說明 backbone 的偽量化對 task_head1 的浮點精度不夠友好，所以需要對模型做以下調整，具體操作方案如下：

對部署的 task_head1 和 task_head2 做 finetune，從而使得 task_head1 和 task_head2 去適應 backbone，直到浮點精度符合預期；

固定 backbone，對 task_head1 做 calib 和 QAT（backbone 的 weight 和 scale 不更新）

固定 backbone，對 task_hea2 做 calib 和 QAT（backbone 的 weight 和 scale 不更新）

固定 backbone 的 weight 和 scale 的方式見下文。

若 step2 驗證出的 head 精度符合預期，則使用以下方案：

固定 backbone 的 weight 和 scale，然后分別對 task_head1 和 task_head2 做 calib/qat；

固定 weight 采用 pytorch 的方法，包括固定 bn 和 stop 梯度更新這兩個操作，如下所示：

#固定bn` `model.eval()` `disable grad` `for param in model.parameters():` `    param.requires_grad = False

征程 6 工具鏈中具有多種 Fix scale 的方式，本文將介紹其中的一種。自定義固定weight和activation的激活 scale 的 qconfig，即配置"averaging_constant": 0，如下為自定義的int8 weight和int8/int16 activation固定 scale 的方式：

from horizon_plugin_pytorch.quantization.qconfig import get_default_qconfig
qat_8bit_fixed_weight_16bit_fixed_act_fake_quant_qconfig = (
    get_default_qconfig(
        weight_qkwargs={
            "dtype": qint8,#weight采用int8量化
            "averaging_constant": 0,#averaging_constant 置為 0 以固定 scale
        },
        activation_qkwargs={
            "dtype": qint16,#采activation用int16量化
            "averaging_constant": 0,#averaging_constant 置為 0 以固定 scale
        },
    )
)
qat_8bit_fixed_weight_8bit_fixed_act_fake_quant_qconfig = (
    get_default_qconfig(
        weight_qkwargs={
            "dtype": qint8,
            "averaging_constant": 0,
        },
        activation_qkwargs={
            "dtype": qint8,
            "averaging_constant": 0,
        },
    )
)

本節將會針對上述步驟展開詳盡且全面的圖文闡釋，通過清晰直觀的圖片和詳細準確的文字說明，為您逐步剖析每個步驟的關鍵要點和操作細節。

此步驟的驗證前提是 backbone 已經完成了 calib/qat，并且偽量化精度已經滿足預期，這里建議 backbone 的偽量化精度要達到浮點精度的 90% 以上。

模型改造：在 backbone 的 forward 代碼的輸入端插入 Quanstub，輸出端插入 Dequanstub；

加載權重：在 prepare 之前加載 task_head1 或者 task_head2 的浮點權重， 在 prepare 之后加載 backbone 的 calib/qat 權重，這里要特別注意加載權重的順序；

calibration：配置模型狀態（如下圖），注意這里模型的狀態要配置為 VALIDATION，然后進行偽量化的精度的驗證；

如果某個 head 的精度較差，那么將固定 backbone 的權重，對此 head 的權重進行微調。