头图

01 背景

首先感谢地平线工具链用户手册和官方提供的示例,给了我很大的帮助,特别是代码注释写了很多的知识点,超赞!要是注释能再详细点,就是超超赞了!下面开始正文。

最近想着学 QAT (量化感知训练)玩玩,大体看了一下地平线的用户手册,不说精度调优之类比较复杂的,光一个 QAT 上手,就感觉对我这种小白不是很友好,捣鼓了好久,感觉在用户手册中很多基础概念都没写,不同模块之间的关联性也没有详细地介绍,后来发现在用户手册 量化感知训练(QAT) 简介章节 ,有这么一句话:
image.png

懂了,没用过 Pytorch 的 QAT,直接看手册学起来有点费劲才是正常滴!

那针对只使用过 Pytorch 在服务器上训练过一些分类、检测模型,没接触过 QAT 的小白,又不想读 PyTorch 官方文档,只想简单入个门,怎么办嘞?欢迎看看这篇文章,提供实操代码和运行步骤,如果文章对你有点作用的话,麻烦收藏+点个赞再走~

该文章参考自 征程6 OE3.0.17中对应的示例以及用户手册

02 基础理论知识

深度学习量化通常是指以 int 类型的数据代替浮点 float 类型的数据进行计算和存储,从而减小模型大小,降低带宽需求,理论上, INT8 量化,与常规的 FP32 模型相比,模型大小减少 4 倍,内存带宽需求减少 4 倍。

量化可以分为 PTQ 与 QAT

  • PTQ:Post-training Quantization,训练后量化,指浮点模型训练完成后,基于一些校准数据,直接通过工具自动进行模型量化的过程,相比 QAT,PTQ 更简单一些,这篇文章不介绍 PTQ。
  • QAT:Quantization aware training,量化感知训练,指浮点模型训练完成后,在模型中插入伪量化节点再进行量化训练的过程,大体过程如下图所示,相比 PTQ,QAT 精度更有保障一些,这篇文章介绍 QAT。
    image.png

小白:图中伪量化节点 FakeQuantize node 是什么?有什么作用?

大黑:从命名看,就是假装量化呗,模拟将数据从 float 类型量化为 int 类型,主要作用于网络的权重和激活(节点输出,不是 relu 这种激活函数的意思)。在 QAT 中,通过使用伪量化节点,可以在训练期间优化模型以适应后续的真实量化操作,从而提高量化模型的准确性和性能。一旦模型训练完成后,伪量化节点将被替换为真实的量化操作,以生成最终的量化模型。

小白:插入伪量化节点后需要 Retraining/Funetuning?感觉很浪费资源的样子…

大黑:通常再多训 1/10 浮点阶段训练的轮数就好了,比如浮点阶段训练了 100epoch,QAT 训个 10epoch 就好,为了精度,浪费就浪费点,小问题!

小白:从上面这个图看,感觉 QAT 还挺简单的,其实目前我就只会用 pytorch 搭一个卷积网络,然后去训练,那我要经历哪些阶段才能得到最终上板部署的模型呢?

大黑:整个过程会涉及到以下几个模型:
image.png

在每个阶段,还有一些需要注意的地方,比如…

小白:停停停,先别急,这里面新名词有点多,先帮我捋捋。float_model 和我直接用 pytorch 搭建的有什么不同吗?calib 是什么?qat_model/qat.bc/quantized.bc 这三者还不是一个意思?板端部署 hbm 模型我知道,就是可以在板子上推理的模型对吧?

大黑:这一连串问题问的挺好,我下面逐个简单解释一下。

  • float_model 和我直接用 pytorch 搭建的有什么不同吗?
    这里 float_model 浮点模型,其实就是在 pytorch 搭建的常规网络输入位置插入 QuantStub 节点、输出位置插入 DeQuantstub 节点,用于区分量化和不量化的边界。
    在 PyTorch 中,QuantStub/DequantStub 是一种用于量化的辅助工具,用于标记量化过程中需要量化/反量化的层或操作,前期浮点训练时可以当它不存在,在定点量化时会自动被替换为对应的量化操作。从普遍而又常规意义上说,想让模型某部分量化,每个输入分支都要插入 QuantStub,每个输出分支都要拆入DeQuantStub,别再追问为什么了,问就是甲鱼的臀部——“规定”。
  • calib 是什么?
    calib 是校准 calibration 的缩写,主要作用是确定量化参数,我们知道,合理的初始量化参数能够显著提升模型精度并加快模型的收敛速度。calibration 就是在浮点模型中插入 Observer,使用少量训练数据,在模型 forward 过程中统计各处的数据分布,以确定合理的量化参数的过程。虽然可以不做 Calibration 直接进行 qat 量化训练,但一般来说,校准对量化训练有益无害,所以推荐大家将此步骤作为必选项。
  • qat_model/qat.bc/quantized.bc 这三者还不是一个意思?
    确实不是一个意思。qat_model:一种插入了伪量化节点的伪量化模型,还是 torch 模型,简单理解为:qat_model 是为了量化训练而存在的模型,里面还“流淌”着浮点的参数,通过伪量化节点在模拟量化而已。qat.bc:相比于 qat_model,多了一步查表算子定点化的操作,精度与 qat_model 可能会存在微小的差异。quantized.bc:模型中浮点算子转换成定点算子,浮点参数全部转换为定点参数,这种转换后的模型称之为quantized_model /定点模型 / 量化模型。
  • 板端部署 hbm 模型我知道,就是可以在板子上推理的模型对吧?
    非常对。
    小白:这些模型是如何生成的?通过图中那几个函数?是地平线封装好的,直接用?

大黑:是的。

03 文件准备与程序运行

参考 OE/samples/ai_toolchain/horizon_model_train_sample/plugin_basic 下的示例,主要使用 fx_mode.py 脚本。注意:这儿的 fx_mode 不是传统意义上 pytorch 中量化方式,而是地平线自己开发的 jit-strip 方式,它具有使用简单、debug 方便等优点,对我们基础用户而言,别学太多,all in 就行,反正不会的可以找地平线技术支持。

  • 文件目录

    (plugin) [xxx plugin_basic]$ tree -L 3
    .
    ├── common.py                                # 功能代码
    ├── fx_mode.py                               # 运行脚本
    ├── data
    │   └── cifar-10-batches-py                  # cifar10数据集
    └── model
      └── mobilenetv2
          └── mobilenet_v2-b0353104.pth         # 预训练权重 

    代码运行,建议在地平线提供的 docker 里运行,当然,如果大家自己会配置本地环境的话,也可以不用 docker。

  • 运行过程
# 通过如下命令,查看一些介绍
python3 fx_mode.py --help

# 生成float-checkpoint.ckpt
python3 fx_mode.py --stage=float 
# 生成calib-checkpoint.ckpt   
python3 fx_mode.py --stage=calib
# 生成qat-checkpoint.ckpt    
python3 fx_mode.py --stage=qat
# 使用定点quantized model evaluate一次      
python3 fx_mode.py --stage=int_infer    
# 编译生成model.hbm,并对script_model进行可视化
python3 fx_mode.py --stage=compile

运行完全程,所有产出物文件如下图:

image.png
跑起来很简单,下面再和大家一起看看代码层面的情况。

04 代码详解

该章节参考地平线用户手册以及自己的理解进行介绍,主要是添加了一些中文注释。

4.1 导入必要依赖

之所以写这一节,主要是希望大家可以从注释中,简单了解各个函数的作用,像 torch、os 等对于我而言,特别基础的导入省略没写,全部的依赖可以看提供的代码。其中,horizon_plugin_pytorch 是地平线基于 PyTorch 开发的量化训练工具,可以理解成 numpy 这种库,里面有很多用于量化训练的的依赖,我们直接用就好了。

  • common.py

    # 定义程序需要接收哪些命令行参数,以及这些参数的类型、默认值等信息。import argparse
    # 使用netron可视化qat.bc或quantized.bc
    from hbdk4.compiler import visualize
    # 从url中下载预训练权重
    from torchvision._internally_replaced_utils import load_state_dict_from_url
    # 用CIFAR10数据集,简单快速
    from torchvision.datasets import CIFAR10
    # 模型预训练权重,用于load_state_dict_from_url
    from torchvision.models.mobilenetv2 import MobileNet_V2_Weights
    # 硬件芯片架构,J6-e:nash-e;J6-m:nash-m
    from horizon_plugin_pytorch import get_march
    # 用于设置伪量化状态,有FakeQuantState.QAT用于qat model train,FakeQuantState.VALIDATION用于qat/calib model eval,FakeQuantState.CALIBRATION用于 calib
    from horizon_plugin_pytorch.quantization import FakeQuantState
    # plugin下的hbdk4,用于export qat.bc和推理bc模型
    from horizon_plugin_pytorch.quantization import hbdk4 as hb4
    # 用于设置qat/calib model 伪量化状态,内参包括FakeQuantState
    from horizon_plugin_pytorch.quantization import set_fake_quantize
    fx_mode.py
    from common import *
    # convert:将export之后的qat.bc 转成 quantized.bc
    # save:保存qat.bc 或 quantized.bc
    from hbdk4.compiler import convert, save
    # torch中的一个类,主要用于将量化操作的结果转换回浮点数,也就是对输出数据转换回浮点数
    from torch.quantization import DeQuantStub
    # 导入MobileNetV2类,用来当父类,目的是构建float_model。from torchvision.models.mobilenetv2 import MobileNetV2
    # 硬件芯片架构,J6-e:nash-e;J6-m:nash-m
    from horizon_plugin_pytorch.march import March, set_march
    # 类似于torch中的类QuantStub,用于将输入数据量化,使用plugin中的QuantStub是因为它支持通过参数手动固定 scale
    from horizon_plugin_pytorch.quantization import QuantStub, prepare
    # calib与qat阶段的量化配置
    from horizon_plugin_pytorch.quantization.qconfig_template import (
      default_calibration_qconfig_setter,
      default_qat_qconfig_setter,    # calib精度差时,qat使用这个
      default_qat_fixed_act_qconfig_setter,    # calib精度好时,qat使用这个
    )

4.2 主函数

看了第2节理论知识部分,主函数部分的代码就是严格执行那几个阶段 stage(详见第2节),很 easy,关于内部细节,在后面几个小节挨个介绍。

def main(
    model: nn.Module,
    stage: str,
    data_path: str,
    model_path: str,
    train_batch_size: int,
    eval_batch_size: int,
    epoch_num: int,
    device_id: int = 0,
    compile_opt: int = 0,
):
    # 对应操作几个阶段的模型
    assert stage in ("float", "calib", "qat", "int_infer", "compile", "visualize",)

    device = torch.device(
        "cuda:{}".format(device_id) if device_id >= 0 else "cpu"
    )

    if not os.path.exists(model_path):
        os.makedirs(model_path, exist_ok=True)

    # 浮点训练阶段优化器
    def float_optim_config(model: nn.Module):
        # This is an example to illustrate the usage of QAT training tool, so
        # we do not fine tune the training hyper params to get optimized         # float model accuracy.
        optimizer = torch.optim.Adam(model.parameters(), weight_decay=2e-4)

        return optimizer, None

    # qat训练阶段优化器
    def qat_optim_config(model: nn.Module):
        # QAT training is targeted at fine tuning model params to match the
        # numerical quantization, so the learning rate should not be too large.
        optimizer = torch.optim.SGD(
            model.parameters(), lr=0.0001, weight_decay=2e-4
        )

        return optimizer, None

    default_epoch_num = {
        "float": 20,     
        "qat": 2,       # 通常float训练epoch数量是qat训练epoch数量的10倍
    }

    if stage in ("float", "qat"):
        if epoch_num is None:
            epoch_num = default_epoch_num[stage]

        train(
            model,
            data_path,
            model_path,
            train_batch_size,
            eval_batch_size,
            epoch_num,
            device,
            float_optim_config if stage == "float" else qat_optim_config,             stage, 
       )

    elif stage == "calib":
        calibrate(
            model,
            data_path,
            model_path,
            train_batch_size,
            eval_batch_size,
            device,
        )

    elif stage == "int_infer":
        int_infer(
            model,
            data_path,
            eval_batch_size,
            device,
        )

    elif stage == "compile":
        compile(
            model,
            model_path,
            compile_opt,
        )
    else:
        visualize(model)

if __name__ == "__main__":
    # 一些基础参数,详见main、get_model_fx中的操作
    args = get_args()
    device = torch.device(
        "cuda:{}".format(args.device_id) if args.device_id >= 0 else "cpu"    )
    model = get_model_fx(args.stage, args.model_path, device)
    main(
        model,
        args.stage,
        args.data_path,
        args.model_path,
        args.train_batch_size, 
        args.eval_batch_size,
        args.epoch_num,
        args.device_id,
        compile_opt=args.opt,
    )

4.3 获取不同阶段模型 get_model_fx

Example input

##############################################################################
# Example input for model tracing and hbir exporting.
##############################################################################
example_input = torch.rand(1, 3, 32, 32)

在代码运行时,有个输入参数 stage 必须配置,表示拿到哪个 model 去整后面的事,当 stage 参数传入(“float”, “calib”, “qat”, “int_infer”)中某一个时,会通过如下函数去获取,具体实现过程解读可见代码注释。

# --------------------------------------------------------------------------#
# We define the model convert pipeline to generate model for each stage.
# --------------------------------------------------------------------------#
def get_model_fx(
    stage: str,
    model_path: str,
    device: torch.device,
    march=March.NASH_E,
) -> nn.Module:
    model_kwargs = dict(num_classes=10, width_mult=1.0)
    float_model = FxQATReadyMobileNetV2(**model_kwargs).to(device)

    if stage == "float":
        # We also could use the origin MobileNetV2 model for float training,
        # because modified QAT ready model can load its params seamlessly.
        float_model = MobileNetV2(**model_kwargs).to(
            device
        )  # these lines are optional

        # Load pretrained model (on ImageNet) to speed up float training.         load_pretrain(float_model, model_path)

        return float_model      # float的时候,到这儿就退出了

    # 浮点训练完成后的权重
    float_ckpt_path = os.path.join(model_path, "float-checkpoint.ckpt") 
    assert os.path.exists(float_ckpt_path)
    float_state_dict = torch.load(float_ckpt_path, map_location=device) 

    # A global march indicating the target hardware version must be setted
    # before prepare qat.
    set_march(march)

    float_model.load_state_dict(float_state_dict)

    # We recommand to use the default_xxx_qconfig_setter, it will
    # enable high precision output if possible. 尾部conv自动配置高精度int32输出
    if stage == "calib":
        qconfig_setter = default_calibration_qconfig_setter
    else:
        qconfig_setter = default_qat_qconfig_setter
    # The op fusion is included in `prepare`. 不需要用户自己实现算子融合
    qat_calib_model = prepare(
        # Catch computation graph of eval mode.
        float_model.eval(),
        # Must give example input to apply model tracing and do model check.
        example_input.to(device),
        qconfig_setter,    # 重中之重,具体介绍后续看时间
    )

    # calib stage时,函数到这儿就会返回了
    if stage == "calib":
        return calib_model

    calib_ckpt_path = os.path.join(model_path, "calib-checkpoint.ckpt")
    assert os.path.exists(calib_ckpt_path)
    calib_state_dict = torch.load(calib_ckpt_path, map_location=device)

    if stage == "qat":
        if os.path.exists(calib_ckpt_path):
            calib_state_dict = torch.load(calib_ckpt_path, map_location=device)
            float_state_dict = None    # 没太理解为什么这儿要把float_state_dict置为None
        else:
            calib_state_dict = None

        if calib_state_dict is not None:
            # qat_model加载的是calib_state_dict!!!qat_calib_model.load_state_dict(calib_state_dict)
        return qat_calib_model    # qat阶段到这儿就退出了

    # int_infer and compile
    qat_ckpt_path = os.path.join(model_path, "qat-checkpoint.ckpt")
    # qat_ckpt_path 或 calib_ckpt_path 有一个就行
    if os.path.exists(qat_ckpt_path):
        qat_calib_model.load_state_dict(
            torch.load(qat_ckpt_path, map_location=device)
        )
    elif os.path.exists(calib_ckpt_path):
        qat_calib_model.load_state_dict(
            torch.load(calib_ckpt_path, map_location=device)
        )
    else:
        raise FileNotFoundError(
            "Do not find saved calib_model or qat_model ckpt "
            "to do int inference."
        )

    qat_ckpt_path = os.path.join(model_path, "qat-checkpoint.ckpt")
    assert os.path.exists(qat_ckpt_path)
    qat_model.load_state_dict(torch.load(qat_ckpt_path, map_location=device))

    # 将 calib/qat_model 模型export导出成伪量化bc、convert转换为定点状态
    # If model has multi inputs, pass them as a tuple.
    qat_hbir_model = hb4.export(qat_calib_model, example_input.to(device))
    quantized_hbir_model = convert(qat_hbir_model, march)
    # 保存两个bc文件,可用于后续直接load,也可以提供给地平线技术支持用于分析
    save(qat_hbir_model, os.path.join(model_path, "qat_hbir_model.bc"))
    save(quantized_hbir_model, os.path.join(model_path, "quantized_hbir_model.bc"))

    return quantized_hbir_model    # int_infer阶段会到这儿才退出

4.4 构建 float_model

从 torchvision.models 中继承 MobileNetV2,微调一下,以支持量化相关操作。模型改造必要的操作有:

  • 在模型所有输入分支前插入 QuantStub
  • 在模型所有输出分支后插入 DequantStub

这部分具体实现过程解读可见代码注释。

# ----------------------------------------------------------------------------#
# Do necessary modify to the MobilenetV2 model from torchvision.
# For FX mode, we need to:
# Insert QuantStub before first layer and DequantStub after last layer.
# Operation replacement and fusion will be carried out automatically (^_^).
# ----------------------------------------------------------------------------#
# 在PyTorch中,QuantStub/DequantStub 是一种用于量化的辅助工具,
# 用于标记量化过程中需要量化/反量化的层或操作,
# 前期浮点训练时当它不存在,在量化时会自动被替换为对应的量化操作
# ----------------------------------------------------------------------------#
# 从torchvision.models中继承MobileNetV2,微调一下
class FxQATReadyMobileNetV2(MobileNetV2):
    def __init__(
        self,
        num_classes: int = 10,      # 实例变量,使用self.来引用变量
        width_mult: float = 0.5,
        inverted_residual_setting: Optional[List[List[int]]] = None,
        round_nearest: int = 8,
    ):
        super().__init__(   # 类变量,使用类名来引用变量,如ClassName.variable_name
            num_classes, width_mult, inverted_residual_setting, round_nearest
        )
        # --------------------------------------------------------------------#
        # 简单理解,在模型首尾部包一层类似于量化反量化操作,每个输入分支都需要包一下
        # --------------------------------------------------------------------#
        # 地平线plugin中的QuantStub可以配置scale
        # 这里的scale=1/128是后面模型输入配置为pyramid必备的
        # pyramid是地平线的芯片上的一个硬件,数据输入可以从这儿来,也可以从DDR来
        # --------------------------------------------------------------------#
        self.quant = QuantStub(scale=1 / 128)   
        self.dequant = DeQuantStub()

    def forward(self, x: Tensor) -> Tensor:
        x = self.quant(x)
        x = super().forward(x)
        x = self.dequant(x)

        return x

4.5 定义常规模型训练与验证的函数

具体实现,看 py 代码就行,很常规。

# --------------------------------------------------------------------------#
# We define dataloaders and other helper functions used in training and evaluation.
# --------------------------------------------------------------------------#

def prepare_data_loaders(
    data_path: str, train_batch_size: int, eval_batch_size: int
) -> Tuple[data.DataLoader, data.DataLoader]:


class AverageMeter(object):
    """Computes and stores the average and current value"""


def accuracy(output: Tensor, target: Tensor, topk=(1,)) -> List[Tensor]:
    """Computes the accuracy over the k top predictions for the specified values of k"""


def train_one_epoch(
    model: nn.Module,
    criterion: Callable,
    optimizer: torch.optim.Optimizer,
    scheduler: Optional[torch.optim.lr_scheduler._LRScheduler],
    data_loader: data.DataLoader,
    device: torch.device,
) -> None:


def evaluate(
    model: nn.Module, data_loader: data.DataLoader, device: torch.device
) -> Tuple[AverageMeter, AverageMeter]:

4.6 float与qat训练代码解读—float_model/qat_model

针对 float_model 和 qat_model 的参数训练,代码解读如下,

# Float and qat share the same training procedure.
def train(
    model: nn.Module,
    data_path: str,
    model_path: str,
    train_batch_size: int,
    eval_batch_size: int,
    epoch_num: int,
    device: torch.device,
    optim_config: Callable,
    stage: str,
):
    train_data_loader, eval_data_loader = prepare_data_loaders(
        data_path, train_batch_size, eval_batch_size
    )

    optimizer, scheduler = optim_config(model)

    best_acc = 0

    for nepoch in range(epoch_num):
        # Train/Eval state must be setted correctly before `set_fake_quantize`
        model.train()
        # --------------------------------------------#
        #   qat模型训练和普通浮点模型训练的不同之处!       # --------------------------------------------#
        if stage == "qat":
            set_fake_quantize(model, FakeQuantState.QAT)

        train_one_epoch(
            model,
            nn.CrossEntropyLoss(),
            optimizer,
            scheduler,
            train_data_loader,
            device,
        )

        model.eval()
        # --------------------------------------------#
        #   qat模型验证精度 和普通浮点模型验证精度 的不同之处!       # --------------------------------------------#
        if stage == "qat":
            set_fake_quantize(model, FakeQuantState.VALIDATION)

        top1, top5 = evaluate(
            model,
            eval_data_loader,
            device,
        )
        print(
            "{} Epoch {}: evaluation Acc@1 {:.3f} Acc@5 {:.3f}".format(
                stage.capitalize(), nepoch, top1.avg, top5.avg
            )
        )

        if top1.avg > best_acc:
            best_acc = top1.avg

            torch.save(
                model.state_dict(),
                os.path.join(model_path, "{}-checkpoint.ckpt".format(stage)),
            )   # 可用于保存 float-checkpoint.ckpt 和 qat-checkpoint.ckpt

    # ----------------------------------------------#
    #   当传入epoch_num=0,用于qat eval
    # ----------------------------------------------#
    if nepoch == 0:
        model.eval()
        if stage == "qat":
            set_fake_quantize(model, FakeQuantState.VALIDATION)

        top1, top5 = evaluate(
            model,
            eval_data_loader,
            device,
        )
        print(
            "{} eval only: evaluation Acc@1 {:.3f} Acc@5 {:.3f}".format(
                stage.capitalize(), top1.avg, top5.avg
            )
        )   # stage.capitalize()表示将字符串首字母大写

    print("Best Acc@1 {:.3f}".format(best_acc))

    return model

4.7 模型校的代码解读—calib_model

float 模型训练完成后,需要进行参数校准,得到 calib_model,如果 calib_model精度满足要求,qat 训练就不需要了,即使 calib_model 精度不行,calib_model_state_dict(校准后的权重)对 qat 训练收敛也非常有帮助。

def calibrate(
    calib_model,
    data_path,
    model_path,
    calib_batch_size,
    eval_batch_size,
    device,
    num_examples=float("inf"),    # float("inf")表示无穷大,主要用于控制使用多少数据进行校准,默认使用所有数据集
):
    # Please note that calibration need the model in eval mode to make BatchNorm act properly.
    calib_model.eval()  # 即使下面用的是train数据集,这儿也是eval
    # set CALIBRATION state will make FakeQuantize in training mode.
    set_fake_quantize(calib_model, FakeQuantState.CALIBRATION)

    train_data_loader, eval_data_loader = prepare_data_loaders(
        data_path, calib_batch_size, eval_batch_size
    )

    with torch.no_grad():
        cnt = 0
        for image, target in train_data_loader:
            image, target = image.to(device), target.to(device)
            calib_model(image)
            print(".", end="", flush=True)
            cnt += image.size(0)
            if cnt >= num_examples:     # 主要用于控制使用多少数据进行校准,默认使用所有数据集
                break
        print()

    # Must set eval mode again before validation, because
    # set CALIBRATION state will make FakeQuantize in training mode.
    calib_model.eval()  
    set_fake_quantize(calib_model, FakeQuantState.VALIDATION)

    top1, top5 = evaluate(
        calib_model,
        eval_data_loader,
        device,
    )
    print(
        "Calibration: evaluation Acc@1 {:.3f} Acc@5 {:.3f}".format(
            top1.avg, top5.avg
        )
    )

    torch.save(
        calib_model.state_dict(),
        os.path.join(model_path, "calib-checkpoint.ckpt"),
    )

    return calib_model

4.8 定点模型评测精度 代码解读—quantized_model

定点模型/quantized 模型/量化模型 eval 推理一下看看精度

def int_infer(
    quantized_model,
    data_path,
    eval_batch_size,
    device,
):
    # hbir do not support dynamic batch size or cuda
    eval_batch_size = 1
    device = torch.device("cpu")

    quantized_model = HbirModule(quantized_model)
    
    _, eval_data_loader = prepare_data_loaders(
        data_path, eval_batch_size, eval_batch_size
    )

    top1, top5 = evaluate(
        quantized_model,
        eval_data_loader,
        device,
    )
    print(
        "Quantized: evaluation Acc@1 {:.3f} Acc@5 {:.3f}".format(
            top1.avg, top5.avg
        )
    )

    return quantized_model

bc/hbir 模型如何推理呢?使用 HbirModule 即可

##############################################################################
# Hbir module wrapper to deal with in-out structures.
##############################################################################
class HbirModule(nn.Module):
    def __init__(self, hbir_module) -> None:
        super().__init__()
        self.hbir_module = hbir_module
        self.input_flattener = hb4.get_hbir_input_flattener(hbir_module)
        self.output_unflattener = hb4.get_hbir_output_unflattener(hbir_module)
    # hbir do not support dynamic batch size or cuda
    def forward(self, input):    # 输入给tuple即可
        flat_input = self.input_flattener(input)
        hbir_rets = self.hbir_module[0](*flat_input)
        rets = self.output_unflattener(hbir_rets)

        return rets

4.9 编译生成上板模型—model.hbm

编译生成上板模型 model.hbm,同时针对 hbm 模型预估性能

def compile(
    quantized_model,
    model_path,
    compile_opt=0,
):
    from hbdk4.compiler import compile, hbm_perf

    hbm_path = os.path.join(model_path, "model.hbm")
    perf_dir = os.path.join(model_path, "perf")

    compile(quantized_model, hbm_path, get_march(), compile_opt)    # 编译生成hbm模型
    hbm_perf(hbm_path, perf_dir)    # 预估hbm模型性能

5. 后续计划

之前写过一篇 X征程3/征程5 零基础学习地平线 QAT 量化感知训练,在 CSDN 上的反馈还行,为了方便直接入手 征程6 的用户,特意写了这篇文章。关于后续,应该会基于 征程6 再学习一些工具链的知识,例如:

  • prepare 阶段的 qconfig_setter 是重中之重,本文并未介绍,后续再看
  • 这个模型比较简单,针对 forward 中有动态循环的怎么处理?不需要处理?后续再看
  • 精度调优过程是什么?精度调优工具怎么用?后续再看

地平线智驾开发者
1 声望2 粉丝

地平线智能驾驶开发者社区旨在连接智能驾驶领域的开发者和对相关技术感兴趣的其他行业开发者、从业者。