嵌入式大赛-海思赛道训练+部署记录(mmpose)

嵌入式大赛-海思赛道训练+部署记录(mmpose)

目录:

  • 训练
  • 模型转换
  • 部署

本项目采用的是商汤MMPose下的YOLOX-Pose

海思开源仓库路径

AIOT/11482_cardio_guidance · HiSpark/2023_hisilicon_embedded_competition - 码云 - 开源中国 (gitee.com)

训练篇

首先,能不训练就不训练,所以最开始,我的目标是找到一个预训练好的网络模型,使用开源的权重参数。但是海思平台是基于caffe1框架的,所以不少现在流行的网络不能直接部署。

期间尝试了Lite-hrnet,但是转换时有问题,结构不适合,所以最后投向了yolox-pose(说到yolox-pose,其实是RM赛场上装甲板的识别给到的灵感)。

通过查阅资料等,可以知道3516不支持Focus层,所以对于yolox网络需要替换一下开头的结构重新进行训练。

所以对于yolox-pose,进需要替换开头的Focus层即可,如下:

def build_stem_layer(self) -> nn.Module:
    """Build a stem layer."""
    # return Focus(
    #     3,
    #     make_divisible(64, self.widen_factor),
    #     kernel_size=3,
    #     norm_cfg=self.norm_cfg,
    #     act_cfg=self.act_cfg)
    return ConvModule(
        3,
        make_divisible(64, self.widen_factor),
        kernel_size=3,
        stride=2,
        padding=1,
        norm_cfg=self.norm_cfg,
        act_cfg=self.act_cfg)

因为要部署在嵌入式设备上,要考虑到设备的性能,所以采用了yolox-pose的nano版本,MMPose官方并未提供,所以从YOLOX的目标检测那里找到了Nano版本的参数,将网络的深度、宽度修改成如下两个系数:

self.depth = 0.33
self.width = 0.25

模型转换篇

首先,需要我们获取到带权重的网络。

注意将下面代码中的姿态检测的路径替换成yolox-pose(也可以尝试其他网络)

from mmpose.apis import MMPoseInferencer

img_path = '/home/nicta100-s55/ai/mmpose/tests/data/coco/000000197388.jpg'

# create the inferencer using the model alias
inferencer = MMPoseInferencer(
    pose2d='/home/nicta100-s55/ai/Lite-HRNet-hrnet/configs/top_down/naive_litehrnet/coco/naive_litehrnet_18_coco_256x192.py',
    pose2d_weights='/home/nicta100-s55/ai/naive_litehrnet_18_coco_256x192.pth'
)

import torch

net = inferencer.model.to('cpu')

直接这样会报错,因为没有.model。

因为采用了MMPose,所以基本没办法自己写网络,然后加载权重到网络中,所以最后采用了一个办法,就是在MMPose中增加了一个函数,get_model,如下:

    def get_model(self):
        return self.model
    if pose2d is not None:
        self.inferencers['pose2d'] = Pose2DInferencer(
            pose2d, pose2d_weights, device, scope, det_model, det_weights,
            det_cat_ids, output_heatmaps)
        self.model = self.inferencers['pose2d'].get_model()

海思官方提供了转换工具。(尝试了,但是我自己网络转换失败了,所以github上又找了一个工具)

pytorch转caffe使用的是brocolli工具(inisis/brocolli: Everything in Torch Fx (github.com))

尤其是提供的QQ交流群,群主非常热心。

我使用时该转换工具有一个要求:就是不能够在forward时输入多个参数,转换时,可能需要在报错的地方,主要是forward的地方进行简单修改,将多余的参数注释,举例如下:

   def forward(self,
                inputs: torch.Tensor
#                data_samples: OptSampleList = None,     # 转的时候改这里了
#                mode: str = 'tensor'
              )-> ForwardResults:
        """The unified entry for a forward process in both training and test.

        The method should accept three modes: "tensor", "predict" and "loss":

        - "tensor": Forward the whole network and return tensor or tuple of
        tensor without any post-processing, same as a common nn.Module.
        - "predict": Forward and return the predictions, which are fully
        processed to a list of :obj:`DetDataSample`.
        - "loss": Forward and return a dict of losses according to the given
        inputs and data samples.

        Note that this method doesn't handle either back propagation or
        parameter update, which are supposed to be done in :meth:`train_step`.

        Args:
            inputs (torch.Tensor): The input tensor with shape
                (N, C, ...) in general.
            data_samples (list[:obj:`DetDataSample`], optional): A batch of
                data samples that contain annotations and predictions.
                Defaults to None.
            mode (str): Return what kind of value. Defaults to 'tensor'.

        Returns:
            The return type depends on ``mode``.

            - If ``mode="tensor"``, return a tensor or a tuple of tensor.
            - If ``mode="predict"``, return a list of :obj:`DetDataSample`.
            - If ``mode="loss"``, return a dict of tensor.
        """
        return self._forward(inputs)
        # if mode == 'loss':
        #     return self.loss(inputs, data_samples)
        # elif mode == 'predict':
        #     return self.predict(inputs, data_samples)
        # elif mode == 'tensor':
        #     return self._forward(inputs, data_samples)
        # else:
        #     raise RuntimeError(f'Invalid mode "{mode}". '
        #                        'Only supports loss, predict and tensor mode')

然后参考下面的网络进行转换即可

import torchvision.models as models
from brocolli.converter.pytorch_caffe_parser import PytorchCaffeParser

net = models.alexnet(pretrained=False)
x = torch.rand(1, 3, 224, 224)
pytorch_parser = PytorchCaffeParser(net, x)
pytorch_parser.convert()
pytorch_parser.save('alexnet')

部署篇

海思3516得到的图像是YUV格式,但是使用的数据集是BGR格式(COCO数据集)。

在caffe转wk格式时,需要选择合适转换。

参考本人当时的提问:海思社区 (hisilicon.com)

这里是当初的板端推理代码,大多数是参考YOLOV2和YOLOV3进行的修改。

AIOT/11482_cardio_guidance/ai_sample/ai_infer_process/ai_infer_process.c · HiSpark/2023_hisilicon_embedded_competition - 码云 - 开源中国 (gitee.com)

尤其注意对齐,网络的三个尺度52,26,13,所以最后梯度是52,28,16。

其余部分参考代码中的带yolox的部分函数代码即可。

你可能感兴趣的:(计算机视觉,海思3516)