7.5.4. 模型分段部署

7.5.4.1. 场景

在一些场景下,用户可能存在将作为一个整体训练的模型拆分成多段进行上板部署的需求。例如对于下图的两阶段检测模型,若 DPP 需要在 CPU 上执行,DPP 的输出(roi)作为 RoiAlign 的输入,则用户需要按虚线框的标注将模型拆分为 Stage1 和 Stage2,分开编译上板。上板运行时,backbone 输出的定点数据直接作为 RoiAlign 的输入。

../../../_images/segmented_deploy.svg

7.5.4.2. 方法

../../../_images/segmented_deploy_method.svg
  1. 模型修改:如上图所示,在正常可以进行量化感知训练的模型基础上,用户需要在 prepare_qat 前在模型分段的分界点后插入 QuantStub,注意若使用了 horizon_plugin_pytorch.quantization.QuantStub,必须设置 scale = None。

  2. QAT 训练:正常作为一个整体对修改后的模型进行量化感知训练,插入的 QuantStub 会将 Stage2 模型输入数据的 scale 记录在 buffer 中

  3. 转定点:正常作为一个整体使用 convert 接口将训练好的 QAT 模型转为定点

  4. 拆分和编译:将模型按照上板后的形态进行拆分,对拆分出的多段模型分别进行 trace 和编译。需要注意的是,虽然在训练时 Stage2 的输入为量化数据,但是在对 Stage2 做 trace 时的 example_input 依然需要是浮点的形式,Stage2 中插入的 QuantStub 会负责给数据配置正确的 scale 并进行量化。