PyTorch 模型快速入门

本章节中，我们将以 torchvision 中的 MobileNetV2 模型为例，介绍 PyTorch 模型从浮点模型改造、Calibration、HBIR 导出、定点模型验证到上板部署的基本流程，帮助您先把一条最小链路跑通。

注意

请先确认您已经完成开发机和开发板环境安装，并且当前环境中可以正常使用 PyTorch、horizon_plugin_pytorch、hbdk4 和相关工具链依赖。

为了尽快展示完整流程，这里使用 CIFAR-10 数据集，而不是 ImageNet-1K。

在开始之前，建议先在开发机侧新建一个 Python 示例文件，例如 quick_start_mobilenetv2.py。后面的 Python 代码可以按章节顺序整理到这个文件中，并在开发机侧通过 python3 quick_start_mobilenetv2.py 运行。

基础准备

下面的示例会直接使用 torchvision.models.MobileNetV2 构建浮点模型，并在模型输入前插入 QuantStub、模型输出后插入 DeQuantStub，使模型具备后续量化和部署所需的基本结构。

导入依赖

这部分依赖覆盖了数据准备、量化准备、HBIR 导出和编译所需的核心模块，后面的代码都直接基于这里的导入展开。

先把下面这部分代码写入 quick_start_mobilenetv2.py 文件开头。

import os
from typing import Callable, List, Optional, Tuple

import numpy as np
import torch
import torch.nn as nn
import torchvision.transforms as transforms
from torch import Tensor
from torch.quantization import DeQuantStub
from torch.utils import data
from torchvision.datasets import CIFAR10
from torchvision.models.mobilenetv2 import MobileNetV2

from horizon_plugin_pytorch.functional import rgb2centered_yuv
from horizon_plugin_pytorch.march import set_march
from horizon_plugin_pytorch.quantization import (
    FakeQuantState,
    PrepareMethod,
    QuantStub,
    get_qconfig,
    prepare,
    set_fake_quantize,
)
from horizon_plugin_pytorch.quantization.hbdk4 import export
from horizon_plugin_pytorch.quantization.qconfig_setter import *
from hbdk4 import compiler as hb4
from hbdk4.compiler import save

基本参数

这里先把示例运行过程中会用到的路径、batch size、训练轮数、目标平台和导出输入统一定义好，后面各阶段直接复用这些参数。

这部分代码可以直接接在前面的 import 后面。

model_path = "./tmp_models/mobilenetv2"
data_path = "./tmp_data/cifar10"
train_batch_size = 256
eval_batch_size = 256
float_epoch_num = 10
qat_epoch_num = 3
march = "nash-e"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Calibration、Export 和开发机定点验证都使用 batch_size=1 的输入
example_input = torch.rand(1, 3, 32, 32, device=device)

os.makedirs(model_path, exist_ok=True)

数据准备

这里直接使用 CIFAR-10，训练侧输入会先转换为 centered_yuv444，与部署链路保持一致。

这一步的输出是两个 DataLoader：一个用于浮点训练和 QAT，另一个用于 Calibration、精度验证和开发机定点验证。

继续在同一个 Python 文件中补齐数据准备函数。

def prepare_data_loaders(
    data_path: str, train_batch_size: int, eval_batch_size: int
) -> Tuple[data.DataLoader, data.DataLoader]:
    normalize = transforms.Normalize(mean=0.0, std=128.0)

    def collate_fn(batch):
        batched_img = torch.stack(
            [
                torch.from_numpy(np.array(example[0], np.uint8, copy=True))
                for example in batch
            ]
        ).permute(0, 3, 1, 2)
        batched_target = torch.tensor([example[1] for example in batch])

        batched_img = rgb2centered_yuv(batched_img)
        batched_img = normalize(batched_img.float())
        return batched_img, batched_target

    train_dataset = CIFAR10(
        data_path,
        train=True,
        transform=transforms.Compose(
            [
                transforms.RandomHorizontalFlip(),
                transforms.RandAugment(),
            ]
        ),
        download=True,
    )
    eval_dataset = CIFAR10(
        data_path,
        train=False,
        download=True,
    )

    train_data_loader = data.DataLoader(
        train_dataset,
        batch_size=train_batch_size,
        sampler=data.RandomSampler(train_dataset),
        num_workers=8,
        pin_memory=True,
        collate_fn=collate_fn,
    )
    eval_data_loader = data.DataLoader(
        eval_dataset,
        batch_size=eval_batch_size,
        sampler=data.SequentialSampler(eval_dataset),
        num_workers=8,
        pin_memory=True,
        collate_fn=collate_fn,
    )

    return train_data_loader, eval_data_loader

评估与训练辅助函数

下面这几个辅助函数分别负责统计精度、执行验证和完成单个 epoch 的训练循环。后面浮点训练、Calibration 后验证和 QAT 都会复用它们。

同样把这部分函数继续追加到 quick_start_mobilenetv2.py 中。

class AverageMeter(object):
    def __init__(self, name: str, fmt=":f"):
        self.name = name
        self.fmt = fmt
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count


def accuracy(output: Tensor, target: Tensor, topk=(1,)) -> List[Tensor]:
    with torch.no_grad():
        maxk = max(topk)
        batch_size = target.size(0)
        _, pred = output.topk(maxk, 1, True, True)
        pred = pred.t()
        correct = pred.eq(target.view(1, -1).expand_as(pred))

        res = []
        for k in topk:
            correct_k = correct[:k].float().sum()
            res.append(correct_k.mul_(100.0 / batch_size))
        return res


def evaluate(
    model: nn.Module, data_loader: data.DataLoader, device: torch.device
) -> Tuple[AverageMeter, AverageMeter]:
    top1 = AverageMeter("Acc@1", ":6.2f")
    top5 = AverageMeter("Acc@5", ":6.2f")

    model.eval()
    with torch.no_grad():
        for image, target in data_loader:
            image, target = image.to(device), target.to(device)
            output = model(image)
            acc1, acc5 = accuracy(output, target, topk=(1, 5))
            top1.update(acc1, image.size(0))
            top5.update(acc5, image.size(0))

    return top1, top5


def train_one_epoch(
    model: nn.Module,
    criterion: Callable,
    optimizer: torch.optim.Optimizer,
    data_loader: data.DataLoader,
    device: torch.device,
) -> None:
    model.train()
    for image, target in data_loader:
        image, target = image.to(device), target.to(device)
        output = model(image)
        loss = criterion(output, target)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

构建浮点模型

普通浮点 MobileNetV2 还不能直接进入后续量化和部署流程，需要先补齐量化入口和出口，在模型输入前插入 QuantStub，在模型输出后插入 DeQuantStub。

class QATReadyMobileNetV2(MobileNetV2):
    def __init__(
        self,
        num_classes: int = 10,
        width_mult: float = 1.0,
        inverted_residual_setting: Optional[List[List[int]]] = None,
        round_nearest: int = 8,
    ):
        super().__init__(
            num_classes, width_mult, inverted_residual_setting, round_nearest
        )
        self.quant = QuantStub()
        self.dequant = DeQuantStub()

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

改造后的模型可以无缝加载改造前模型的参数，因此若已有训练好的浮点模型，直接加载即可，否则需要正常进行浮点训练。

下面这段代码完成浮点训练，并在 model_path 下保存表现最好的浮点 checkpoint。后续 Calibration 会直接从这里产出的浮点参数继续。

把这部分模型定义和训练入口继续写到同一个 Python 文件中。等后面的 Calibration、QAT、导出和编译逻辑也补齐后，再统一在开发机侧执行这个脚本。

float_model = QATReadyMobileNetV2().to(device)
train_data_loader, eval_data_loader = prepare_data_loaders(
    data_path, train_batch_size, eval_batch_size
)

optimizer = torch.optim.Adam(
    float_model.parameters(), lr=0.001, weight_decay=1e-3
)
best_acc = 0

for nepoch in range(float_epoch_num):
    train_one_epoch(
        float_model,
        nn.CrossEntropyLoss(),
        optimizer,
        train_data_loader,
        device,
    )

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

    if top1.avg > best_acc:
        best_acc = top1.avg
        torch.save(
            float_model.state_dict(),
            os.path.join(model_path, "float-checkpoint.ckpt"),
        )

如果您已经有浮点 checkpoint，可以直接用下面这段代码替换前面的浮点训练部分，然后继续执行后面的 Calibration。

float_model = QATReadyMobileNetV2().to(device)
float_model.load_state_dict(
    torch.load(os.path.join(model_path, "float-checkpoint.ckpt"), map_location=device)
)

Calibration

模型改造完成并准备好浮点参数后，下一步就可以做 Calibration。

这一步会在模型中插入 Observer，统计各处数据分布，并据此生成一组可用于部署的初始量化参数。对这个 MobileNetV2 示例来说，Calibration 后即可先继续验证部署链路是否正常。

这里我们使用 QconfigSetter 方式来配置 qconfig。

set_march(march)

calib_model = prepare(
    float_model,
    example_inputs=example_input,
    qconfig_setter=QconfigSetter(
        reference_qconfig=get_qconfig(),
        templates=[
            ModuleNameTemplate({"": qint8}),
            ConvDtypeTemplate(),
            MatmulDtypeTemplate(),
        ],
    ),
)

calib_model.eval()
set_fake_quantize(calib_model, FakeQuantState.CALIBRATION)
with torch.no_grad():
    for image, _ in eval_data_loader:
        image = image.to(device)
        calib_model(image)

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

Calibration 完成后，建议再验证一次伪量化精度：

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
    )
)

Calibration 完成后，可以先判断当前量化结果是否满足需求，如果还需要进一步提升精度，再继续执行下面的 QAT。也就是说，Calibration 在这里既负责生成初始量化参数，也负责帮您判断是否还需要继续做 QAT。

量化感知训练

如果 Calibration 后的结果仍然不满足需求，请您进一步对模型进行量化感知训练，即 QAT。这一步会加载 Calibration 产出的量化参数，在训练过程中继续微调模型权重，以进一步提升量化后的精度。

这里，我们同样沿用当前推荐的 QconfigSetter + PrepareMethod.JIT_STRIP 组合，与前面的 Calibration 保持一致。

qat_model = prepare(
    float_model,
    example_inputs=example_input,
    qconfig_setter=QconfigSetter(
        reference_qconfig=get_qconfig(),
        templates=[
            ModuleNameTemplate({"": qint8}),
            ConvDtypeTemplate(),
            MatmulDtypeTemplate(),
        ],
    ),
)
qat_model.load_state_dict(calib_model.state_dict())

optimizer = torch.optim.Adam(
    qat_model.parameters(), lr=1e-3, weight_decay=1e-4
)
best_acc = 0

for nepoch in range(qat_epoch_num):
    qat_model.train()
    set_fake_quantize(qat_model, FakeQuantState.QAT)

    train_one_epoch(
        qat_model,
        nn.CrossEntropyLoss(),
        optimizer,
        train_data_loader,
        device,
    )

    qat_model.eval()
    set_fake_quantize(qat_model, FakeQuantState.VALIDATION)
    top1, top5 = evaluate(qat_model, eval_data_loader, device)
    print(
        "QAT Epoch {}: evaluation Acc@1 {:.3f} Acc@5 {:.3f}".format(
            nepoch, top1.avg, top5.avg
        )
    )

    if top1.avg > best_acc:
        best_acc = top1.avg
        torch.save(
            qat_model.state_dict(),
            os.path.join(model_path, "qat-checkpoint.ckpt"),
        )

定点模型验证与编译

当伪量化结果满足需求后，下一步先导出 HBIR 模型，再转为定点模型并完成编译。

导出 HBIR 模型

导出时会根据当前部署逻辑执行一次完整 forward，并将执行到的算子逐个转为 HBIR 节点。

这一步的输出是 deploy.bc，它仍然属于导出阶段产物，后面的 Convert 会以它作为输入。

注解

下方示例代码默认导出进行了 QAT 的模型，若未进行 QAT，可将示例代码中的 qat_model 替换为 calib_model。

deploy_model = qat_model

deploy_model.eval()
set_fake_quantize(deploy_model, FakeQuantState.VALIDATION)
hbir_model = export(deploy_model, (example_input,))
save(hbir_model, os.path.join(model_path, "deploy.bc"))

转定点模型

得到导出的 HBIR 模型后，便可将模型转为定点模型：

这一步会把 deploy.bc 转为真正的定点模型 quantized.bc。通常在进入编译之前，先看这一阶段的精度会更稳妥。

hbir_quantized_model = hb4.convert(hbir_model, march)
save(hbir_quantized_model, os.path.join(model_path, "quantized.bc"))

开发机结果确认

转定点后，可以先在开发机侧验证 quantized.bc 对应的定点结果是否正常：

_, eval_hbir_data_loader = prepare_data_loaders(data_path, train_batch_size, 1)


def evaluate_hbir(
    model: hb4.Module, data_loader: data.DataLoader
) -> Tuple[AverageMeter, AverageMeter]:
    top1 = AverageMeter("Acc@1", ":6.2f")
    top5 = AverageMeter("Acc@5", ":6.2f")

    for image, target in data_loader:
        image = torch.as_tensor(image.cpu().numpy())
        target = target.cpu()
        output = model["forward"].feed({"_input_0": image})["_output_0"]
        acc1, acc5 = accuracy(output, target, topk=(1, 5))
        top1.update(acc1, image.size(0))
        top5.update(acc5, image.size(0))

    return top1, top5


top1, top5 = evaluate_hbir(hbir_quantized_model, eval_hbir_data_loader)
print(
    "Quantized model: evaluation Acc@1 {:.3f} Acc@5 {:.3f}".format(
        top1.avg, top5.avg
    )
)

编译模型

确认开发机侧定点结果正常后，再继续编译生成 model.hbm：

编译完成后，后面的开发板动态性能验证和 C++ 环境验证都将直接使用这里产出的 model.hbm。

hb4.compile(
    hbir_quantized_model,
    os.path.join(model_path, "model.hbm"),
    march,
    opt=1,
)

hb4.hbm_perf(
    os.path.join(model_path, "model.hbm"),
    output_dir=model_path,
)

完成这一步后，开发机侧已经得到了可上板使用的 model.hbm，后续即可继续开发板动态性能验证和部署验证。

编译完成后，model_path 下通常会生成：

./tmp_models/mobilenetv2/
├── float-checkpoint.ckpt
├── calib-checkpoint.ckpt
├── qat-checkpoint.ckpt
├── deploy.bc
├── quantized.bc
├── model.hbm
├── model.html
└── model.json

开发板动态性能验证

拿到 model.hbm 后，可以先上板实测模型的动态性能，确认模型已经能在开发板上正常运行：

这一步的目标很简单：先确认模型已经能在开发板上正常加载和执行，再继续后面的 C++ 环境验证。

# 将模型拷贝至开发板
scp ./tmp_models/mobilenetv2/model.hbm root@{board_ip}:/userdata

# 登录开发板
ssh root@{board_ip}
cd /userdata

# 单线程串行评测 latency
hrt_model_exec perf --model_file model.hbm --thread_num 1 --frame_count 1000

# 多线程并发评测 FPS
hrt_model_exec perf --model_file model.hbm --core_id 0 --thread_num 8 --frame_count 1000

如果开发板环境里没有 hrt_model_exec，可以先回到 OE 包的 package/board 路径，重新执行对应的板端安装脚本。

开发板 C++ 环境验证

在开发板端，地平线提供了统一异构计算平台 UCP，来帮助您快速完成模型的部署工作，并提供了相关示例。

如果您希望继续了解模型部署和推理 API 的基础使用方式，可以先阅读基础示例包使用说明，再参考基础示例中的目录结构、构建方式和接口调用顺序继续完成板端验证。

如果您希望继续参考完整示例工程中的性能和精度评测方式，可以再阅读 samples/ucp_tutorial/dnn/ai_benchmark 对应的 AI-Benchmark 使用说明。

当模型的性能和精度验证都符合预期后，即可参考运行时开发继续完成上层应用开发。

#PyTorch 模型快速入门

#基础准备

#导入依赖

#基本参数

#数据准备

#评估与训练辅助函数

#构建浮点模型

#Calibration

#量化感知训练

#定点模型验证与编译

#导出 HBIR 模型

#转定点模型

#开发机结果确认

#编译模型

#开发板动态性能验证

#开发板 C++ 环境验证