PyTorch 框架的 Yolov5 移植_寒武纪yolov5移植

前言

本文对开源yolov5s模型进行寒武纪平台的移植

整个移植过程分为模型结构转换、添加后处理算子框架代码、模型量化、在线推理和离线推理共五个步骤。

1. 模型结构转换

对于原始Yolov5网络的后处理部分的逻辑，Cambricon-PyTorch直接使用一个大的BANGC算子完成后处理的计算，需要对原生的pytorch网络进行修改，将后处理部分的整体计算换成BANGC算子。

具体做法为是将yololayer层替换成了yolov5_detection_output，把三个yololayer的输入传给了yolov5_detection_output。修改部分在yolo.py中，如下：

if x[0].device.type == 'mlu':

         for i in range(self.nl):

             x[i] = self.m[i](x[i])  # conv

        y = x[i].sigmoid()

        output.append(y)

    detect_out = torch.ops.torch_mlu.yolov5_detection_output(output[0], output[1], output[2],

                                                             self.anchors_list,self.nc, self.num_anchors,