轻量级实时跟踪算法NanoTrack在瑞芯微RK3588上的部署以及使用

文章目录

前言

一、模型转换

1.环境配置

2.模型解构

二、rk3588平台使用

1.模型初始化

2.推理

---

---

前言

github: https://github.com/Try2ChangeX/NanoTrack_RK3588_python: python版本基于rk3588的NanoTrack,每秒可达120FPS

主要参考：

SiamTrackers/NanoTrack at master · HonglinChu/SiamTrackers · GitHub

GitHub - rockchip-linux/rknn-toolkit2

        本文主要记录了在rk3588开发板上开发并运行nanotack跟踪算法，主要参考上面大佬的SiamTrackers，然后将模型转为rk3588能使用的.rknn模型并成功运行。主要记录该算法在板子上运行需要注意的一些坑和细节点。

---

一、模型转换

1.环境配置

rk3588对应的python版本的开发库是：GitHub - rockchip-linux/rknn-toolkit2，模型转换需要在linux-x86机器上进行转换，可以参考官方提供的packages去配置相关环境，目前rknn_toolkit2就两个版本，如下。可以选择在ubuntu18.04的机器上python3.6环境下安装第一个package，最好联网安装，会有很多依赖包。

rknn_toolkit2的使用文档可参考./rknn-toolkit2-1.3/doc/下开发文档

由于rk3588板子系统为aarch64架构，需在板子上安装对应的rknn_toolkit_lite2，同样，环境可根据其packages进行配置，这里注意，只支持python3.7和python3.9,需提前配置好板子上的python版本，可选择使用虚拟环境。

 rknn_toolkit_lite2的使用文档可参考./rknn_toolkit_lite2/docs/下开发文档

2.模型解构

关于模型的前处理，后处理等部分，可参考代码，本文不对算法做过多描述。

模型的输入分为两部分：

模板T的input输入为[1,3,127,127]，经过backbone特征提取网络得到output1[1,48,8,8]的模板特征；

图像X的input输入为[1,3,255,255]，经过backbone特征提取网络得到output2[1,48,16,16]的图像特征；

output1和output2会被送入模型head部分，得到两部分输出,输出尺寸分别为[1,48,8,8]和[1,48,16,16]。

考虑到rknn在转换时需提前设置好输出节点，所以这里我将模型拆成了3个部分，分别为：

input	输入数量	model	output	输出数量	备注
模板T输入[1,3,127,127]	1	backbone	[1,48,8,8]	1	仅在init时进行推理
图像X输入[1,3,255,255]	1	backbone	[1,48,16,16]	1	update时一直推理
[1,48,8,8],[1,48,16,16]	2	head	[1,2,16,16] [1,4,16,16]	2	双输入双输出

提前将训练好的模型分别导出为.pt格式，方便后续转.rknn

# backbone 模板特征提取模型
trace_model = torch.jit.trace(backbone_net, torch.Tensor(1, 3, 127, 127))
trace_model.save('./models/pt/backbone_127.pt')

# backbone 图像特征提取模型
trace_model = torch.jit.trace(backbone_net, torch.Tensor(1, 3, 255, 255))
trace_model.save('./models/pt/backbone_255.pt')

# head 模型
trace_model = torch.jit.trace(model, (torch.Tensor(1, 48, 8, 8), torch.Tensor(1, 48, 16, 16)))
trace_model.save('./models/pt/head.pt')

然后通过rknn_toolkit2转为.rknn模型，转换时注意是否需要归一化，以及target_platform，为防止出错，最好直接指定目标平台。如：模板特征提取的模型部分转换如下

model = './pt/nanotrack_backbone_T127.pt'
        
input_size_list = [[1, 3, 127, 127]]

# Create RKNN object
rknn = RKNN(verbose=True)

# Pre-process config
print('--> Config model')
# 注意要将target_platform指定清楚
rknn.config(mean_values=[0, 0, 0], std_values=[1, 1, 1], target_platform='rk3588')
print('done')

# Load model
print('--> Loading model')
ret = rknn.load_pytorch(model=model, input_size_list=input_size_list)
if ret != 0:
    print('Load model failed!')
    exit(ret)
print('done')

# Build model
print('--> Building model')
ret = rknn.build(do_quantization=False)
if ret != 0:
    print('Build model failed!')
    exit(ret)
print('done')

# Export rknn model
print('--> Export rknn model')
ret = rknn.export_rknn('./rknn/backbone_127.rknn')
if ret != 0:
     print('Export rknn model failed!')
     exit(ret)
print('done')

到此，模型转换完成。

可以提前在linux-x86的机器上用用模拟环境运行，测试模型是否转换正常，输出是否符合预期。

模拟环境使用可直接参考./rknn-toolkit2-1.3/examples/下的程序示例。

二、rk3588平台使用

1.模型初始化

仅仅只做模型初始化

self.rknn_Tback = RKNNLite()

# load RKNN model
print('--> Load RKNN model')
ret = self.rknn_Tback.load_rknn(Tback_weight)
if ret != 0:
    print('Load RKNN model failed')
    exit(ret)
print('done')

# init runtime environment
print('--> Init runtime environment')

ret = self.rknn_Tback.init_runtime(core_mask=RKNNLite.NPU_CORE_0)
if ret != 0:
    print('Init runtime environment failed')
    exit(ret)
print('done')

2.推理

需要注意的是，在rknn_lite里面，inference时，数据的输入都是需按照NHWC格式储存，所以，在送进模型之前需要将数据进行换轴操作。（与模拟环境运行不同，模拟环境可以指定data_format='nchw'），这里在rknn_toolkit_lit2的开发文档中说明了。（还是要多看文档，rknn很多组件没有开发完全，使用起来有很多限制）

inference时候的参考代码如下：

back_T_in = z_crop.transpose((0,2,3,1)) //提前换成‘nhwc’格式
self.Toutput = self.rknn_Tback.inference(inputs=[back_T_in])

预测得到的数据都是一个list类型，只有一个输出则list长度为1，有两个输出则list长度为2。

附：

有C++部署需求的小伙伴可参考：

https://github.com/Try2ChangeX/NanoTrack_Libtorch: C++版本基于libtorch的NanoTrack跟踪算法

基于libtorch部署的Nanotrack跟踪算法，已验证跑通。