rknn上部署人体关键点检测模型（2）

模型转换、部署

源代码仓库地址为：https://github.com/HRNet/Lite-HRNet.git，不过更推荐使用开源项目open-mmlab(https://openmmlab.com/),该项目有多种类丰富的开源代码，实现了分类、检测

、识别、定位等多种模型。其中的mmpose项目(https://github.com/open-mmlab/mmpose.git)即实现了lite-hrnet模型，我们可以通过mmpose项目中的脚本，轻松的训练、导出模型。

rknn的部署一般先要将模型转换为onnx模型，然后将onnx模型转换为rknn模型，其中转换为onnx模型的脚本我们可以利用mmpose/tools/deployment/pytorch2onnx.py脚本进行转换。也可以参考源代码通过以下步骤导入，转换模型：

import argparse

import mmcv

import numpy as np

import torch

from mmcv.runner import load_checkpoint

from time import time

from models import build_posenet

from mmcv.onnx.symbolic import register_extra_symbolics

def parse_args():

parser = argparse.ArgumentParser(

description=’Convert models to ONNX’)

parser.add_argument(‘–config’,

default=”configs/top_down/lite_hrnet/coco/litehrnet_30_coco_256x192.py”,

help=’config file path’)

parser.add_argument(‘–checkpoint’,

default=”litehrnet_30_256x192.pth”,

help=’checkpoint file’)

parser.add_argument(‘–output-file’, type=str, default=’lite_hr_30_256.onnx’)

parser.add_argument(

‘–shape’,

type=int,

default=(1, 3, 256, 192),

help=’input size’)

args = parser.parse_args()

return args

if __name__ == ‘__main__’:

args = parse_args()

cfg = mmcv.Config.fromfile(args.config)

# build the model

model = build_posenet(cfg.model)

if hasattr(model, ‘forward_dummy’):

model.forward = model.forward_dummy

else:

raise NotImplementedError(

‘Please implement the forward method for exporting.’)

checkpoint = load_checkpoint(model, args.checkpoint,map_location=’cpu’)

model.eval()

pytorch2onnx(

model,

args.shape,

output_file=args.output_file)

这样model就是导入的模型了，之后使用pytorch2onnx函数即可进行转换。

def pytorch2onnx(model,

input_shape,

output_file=’a.onnx’):

model.eval()

one_img = torch.randn(input_shape)

register_extra_symbolics(11)

torch.onnx.export(

model,

one_img,

output_file,

export_params=True,

do_constant_folding=True,

opset_version=11

)

print(f’Successfully exported ONNX model: {output_file}’)

导出模型过程中可能出现RuntimeError: Failed to export an ONNX attribute ‘onnx::Gather’, since it’s not constant, please try to make things (e.g., kernel size) static if possible问题。官方代码中使用了自适应的池化函数adaptive_avg_pool2d，导致无法正常导出。需要修改官方github中的models\backbones\litehrnet.py，将其中的自适应池化函数改为正常的池化函数，因此需要计算出正常池化函数的几个参数：核大小kernel_size、步长stride_size、填充padding（一般都是0，可以忽略这个参数）。修改其中的CrossResolutionWeighting类的forward。

def forward(self, x):

#mini_size = x[-1].size()[-2:]

#out = [F.adaptive_avg_pool2d(s, mini_size) for s in x[:-1]] + [x[-1]]

stridesize_list, kernelsize_list, output_size = self.get_avgpool_para(x)

out = [torch.nn.AvgPool2d(kernel_size=kernelsize_list.tolist(),stride=stridesize_list.tolist())(s) for s,kernel_size,stride_size in

zip(x[:-1],kernelsize_list,stridesize_list) ] + [x[-1]]

out = torch.cat(out, dim=1)

out = self.conv1(out)

out = self.conv2(out)

out = torch.split(out, self.channels, dim=1)

out = [

s * F.interpolate(a, size=s.size()[-2:], mode=’nearest’)

for s, a in zip(x, out)

]

return out

get_avgpool_para实现如下：

def get_avgpool_para(self,x):

output_size = np.array(x[-1].size()[-2:])

stride_size_list = []

kernel_size_list = []

for index in range(len(x)):

input_size = np.array(x[index].size()[-2:])

stride_size = np.floor(input_size / output_size ).astype(np.int32)

kernel_size = input_size – (output_size – 1) * stride_size

stride_size_list.append(stride_size)

kernel_size_list.append(kernel_size)

return stride_size_list,kernel_size_list,output_size

转换为onnx模型后通过rknn中onnx转rknn脚本进行转换。将onnx模型转换为rknn模型。rknn开发套件地址为:https://github.com/rockchip-linux/rknn-toolkit2.git

RKNN-Toolkit2是一款软件开发套件，供用户在PC和瑞芯微NPU平台（RK3566、RK3568、RK3588、RK3588S、RV1103、RV1106）上进行模型转换、推理和性能评估。RKNN-Toolkit-Lite2为Rockchip NPU平台（RK3566、RK3568、RK3588、RK3588S）提供Python编程接口，帮助用户部署RKNN模型，加速AI应用落地。

我们可以按照rknn-toolkit2/examples/onnx/yolov5/的例子，修改后转换我们已经生成了的onnx模型。

注意事项：

rknn.config中预处理方式与模型预处理方式要相同，即mean_values、std_values要与实际训练过程中预处理方式相同。mean_values：输入的均值。参数格式是一个列表，列表中包含一个或多个均值子列表，多输入模型对应多个子列表，每个子列表的长度与该输入的通道数一致，例如[[128,128,128]]，表示一个输入的三个通道的值减去 128。如果quant_img_RGB2BGR设置成 True，则优先做 RGB2BGR 转换，再做减均值。 std_values：输入的归一化值。参数格式是一个列表，列表中包含一个或多个归一化值子列表，多输入模型对应多个子列表，每个子列表的长度与该输入的通道数一致， 18 例如[[128,128,128]]，表示设置一个输入的三个通道的值减去均值后再除以128。如果 quant_img_RGB2BGR 设置成 True，则优先做 RGB2BGR 转换，再减均值和除以归一化值。

如果要连接rknn设备进行推理测试需要在config及init_runtime中设置target：指定 RKNN 模型是基于哪个目标芯片平台生成的。目前支持“rk3566”、“rk3568”、“rk3588”、“rv1103”、“rv1106”。值得注意的是init_runtime中参数为target而config中参数为target_platform。

转换为rknn模型后可以通过rknn.eval_perf(inputs=[image], is_print=True)评估模型性能，查看每层消耗的时间，针对性修改模型从而提高模型的推理速度，降低输入层尺寸也是提速度最简单的方式之一。

rk芯片可能存在的问题:

（1）在转换过程中，报错有些算子不支持，可以自己修改比较通用的算子后重新转换模型。

（2）算子支持不同步，如在3568上支持的算子在3588上反而不支持。

（3）算子在cpu上执行，有些层的算子在硬件cpu上执行，导致速度很慢。

**************************************宣传分割线**********************************

深圳原数科技提供嵌入式软硬件、服务器、APP、代工生产一站式服务，对国产MCU和SOC（全志，瑞芯微）有丰富开发经验，特别是在低成本的wifi图传和AI视觉识别领域有深厚沉淀，相关产品出货超过50W台。近来顺应市场需求陆续推出人形追踪云台、虚拟背景直播一体机等落地产品，欢迎联系咨询,李生13686461548微信同号，也可官网查询相关信息www.atom-math.com