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



    前言

    github: GitHub - WindyXiang/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输入数量modeloutput输出数量备注
    模板T输入[1,3,127,127]1backbone[1,48,8,8]1仅在init时进行推理
    图像X输入[1,3,127,127]1backbone[1,48,16,16]1update时一直推理
    [1,48,8,8],[1,48,16,16]2head

    [1,2,16,16]

    [1,4,16,16]

    2双输入双输出

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

    1. # backbone 模板特征提取模型
    2. trace_model = torch.jit.trace(backbone_net, torch.Tensor(1, 3, 127, 127))
    3. trace_model.save('./models/pt/backbone_127.pt')
    4. # backbone 图像特征提取模型
    5. trace_model = torch.jit.trace(backbone_net, torch.Tensor(1, 3, 255, 255))
    6. trace_model.save('./models/pt/backbone_255.pt')
    7. # head 模型
    8. trace_model = torch.jit.trace(model, (torch.Tensor(1, 48, 8, 8), torch.Tensor(1, 48, 16, 16)))
    9. trace_model.save('./models/pt/head.pt')

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

    1. model = './pt/nanotrack_backbone_T127.pt'
    2. input_size_list = [[1, 3, 127, 127]]
    3. # Create RKNN object
    4. rknn = RKNN(verbose=True)
    5. # Pre-process config
    6. print('--> Config model')
    7. # 注意要将target_platform指定清楚
    8. rknn.config(mean_values=[0, 0, 0], std_values=[1, 1, 1], target_platform='rk3588')
    9. print('done')
    10. # Load model
    11. print('--> Loading model')
    12. ret = rknn.load_pytorch(model=model, input_size_list=input_size_list)
    13. if ret != 0:
    14. print('Load model failed!')
    15. exit(ret)
    16. print('done')
    17. # Build model
    18. print('--> Building model')
    19. ret = rknn.build(do_quantization=False)
    20. if ret != 0:
    21. print('Build model failed!')
    22. exit(ret)
    23. print('done')
    24. # Export rknn model
    25. print('--> Export rknn model')
    26. ret = rknn.export_rknn('./rknn/backbone_127.rknn')
    27. if ret != 0:
    28. print('Export rknn model failed!')
    29. exit(ret)
    30. print('done')

    到此,模型转换完成。

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

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

    二、rk3588平台使用

    1.模型初始化

    仅仅只做模型初始化

    1. self.rknn_Tback = RKNNLite()
    2. # load RKNN model
    3. print('--> Load RKNN model')
    4. ret = self.rknn_Tback.load_rknn(Tback_weight)
    5. if ret != 0:
    6. print('Load RKNN model failed')
    7. exit(ret)
    8. print('done')
    9. # init runtime environment
    10. print('--> Init runtime environment')
    11. ret = self.rknn_Tback.init_runtime(core_mask=RKNNLite.NPU_CORE_0)
    12. if ret != 0:
    13. print('Init runtime environment failed')
    14. exit(ret)
    15. print('done')
     
    

    2.推理

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

    inference时候的参考代码如下:

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

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

    附:

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

    GitHub - WindyXiang/NanoTrack_Libtorch: C++版本基于libtorch的NanoTrack跟踪算法

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

  • 相关阅读:
    全排列:让我看到未来所有的可能 -> 跨越CV小白的回溯算法
    wps:样式集的使用【笔记】
    向NXP官网Linux内核添加ALPHA开发板
    Java实现3DES加密解密(DESede/ECB/PKCS5Padding使用)
    SALOME源码分析:GUI模块
    优化Redis缓存淘汰机制解决性能测试中报错率逐渐攀升问题
    基于SpringBoot Vue宠物领养系统
    Java项目:SSM游戏点评网站
    Operations for Algorithms
    华为程序员,985本科36岁,被公司解约:中年人路在何方?
  • 原文地址:https://blog.csdn.net/X_windy/article/details/126153227