yolov5 pt-＞onnx-＞om yolov5模型转onnx转om模型转换

yolov5 pt->onnx->om
yolov5-6.1版本
models/yolo.py

Detect函数修改

class Detect(nn.Module):
	def forward(self, x):
	        z = []  # inference output
	        for i in range(self.nl):
	            x[i] = self.m[i](x[i])  # conv
	            bs, _, ny, nx = x[i].shape  # x(bs,255,20,20) to x(bs,3,20,20,85)
	            x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
	            y = x[i].sigmoid()
	            z.append(y.view(bs, -1, self.no))
	
	        return torch.cat(z, 1)

common.py Focus修改(说是可以提升Slice算子性能，待测，可不改)

class Focus(nn.Module):
    def forward(self, x):  # x(b,c,w,h) -> y(b,4c,w/2,h/2)
        # <==== 修改内容
        if torch.onnx.is_in_onnx_export():
            a, b = x[..., ::2, :].transpose(-2, -1), x[..., 1::2, :].transpose(-2, -1)
            c = torch.cat([a[..., ::2, :], b[..., ::2, :], a[..., 1::2, :], b[..., 1::2, :]], 1).transpose(-2, -1)
            return self.conv(c)
        else:
            return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1))
        # return self.conv(self.contract(x))
        # =====>

python3 export.py --weights "/home/yolo_object_detection/yolov5-6.1/runs/train/exp49/weights/best.pt"   --img-size=640 --batch-size 1

转om模型

atc --model="/hsw/best.onnx" --framework=5 --output=/hsw/smoke_2 --soc_version=Ascend310  --input_shape="images:1,3,640,640" --insert_op_conf="/hsw/aipp_nv12.cfg"

aipp_nv12.cfg 内容如下

aipp_op{
aipp_mode:static
input_format : RGB888_U8
csc_switch: false
rbuv_swap_switch: true
var_reci_chn_0 :0.003922
var_reci_chn_1 :0.003922
var_reci_chn_2 :0.003922
}

//om 模型yolov5模型后处理
void yolov5_sm_call::nms(std::vector<BoxInfoCar>& input_boxes)
{
    sort(input_boxes.begin(), input_boxes.end(), [](BoxInfoCar a, BoxInfoCar b) { return a.score > b.score; });
    std::vector<float> vArea(input_boxes.size());
    for (int i = 0; i < int(input_boxes.size()); ++i)
    {
        vArea[i] = (input_boxes.at(i).x2 - input_boxes.at(i).x1 + 1)
            * (input_boxes.at(i).y2 - input_boxes.at(i).y1 + 1);
    }

    std::vector<bool> isSuppressed(input_boxes.size(), false);
    for (int i = 0; i < int(input_boxes.size()); ++i)
    {
        if (isSuppressed[i]) { continue; }
        for (int j = i + 1; j < int(input_boxes.size()); ++j)
        {
            if (isSuppressed[j]) { continue; }
            float xx1 = (std::max)(input_boxes[i].x1, input_boxes[j].x1);
            float yy1 = (std::max)(input_boxes[i].y1, input_boxes[j].y1);
            float xx2 = (std::min)(input_boxes[i].x2, input_boxes[j].x2);
            float yy2 = (std::min)(input_boxes[i].y2, input_boxes[j].y2);

            float w = (std::max)(float(0), xx2 - xx1 + 1);
            float h = (std::max)(float(0), yy2 - yy1 + 1);
            float inter = w * h;
            float ovr = inter / (vArea[i] + vArea[j] - inter);

            if (ovr >= this->nmsThreshold)
            {
                isSuppressed[j] = true;
            }
        }
    }
    // return post_nms;
    int idx_t = 0;
    input_boxes.erase(remove_if(input_boxes.begin(), input_boxes.end(), [&idx_t, &isSuppressed](const BoxInfoCar& f) { return isSuppressed[idx_t++]; }), input_boxes.end());
}


void yolov5_sm_call::detect(float ratioh, float ratiow, void* out_, std::vector<BBox_v4>& result)
{
    int newh = inpWidth, neww = inpHeight, padh = 0, padw = 0;

    std::vector<BoxInfoCar> generate_boxes;
    int n = 0, q = 0, i = 0, j = 0, row_ind = 0, k = 0; ///xmin,ymin,xamx,ymax,box_score, class_score
    const float* pdata = static_cast<float*>(out_);
    for (n = 0; n < this->num_stride; n++)
    {
        const float stride = pow(2, n + 3);
        int num_grid_x = (int)ceil((this->inpWidth / stride));
        int num_grid_y = (int)ceil((this->inpHeight / stride));
        for (q = 0; q < 3; q++)    ///anchor
        {
            const float anchor_w = this->anchors[n * 6 + q * 2];
            const float anchor_h = this->anchors[n * 6 + q * 2 + 1];
            for (i = 0; i < num_grid_y; i++)
            {
                for (j = 0; j < num_grid_x; j++)
                {
                    float box_score = pdata[4];
                    if (box_score > this->objThreshold)
                    {
                        int max_ind = 0;
                        float max_class_socre = 0;
                        for (k = 0; k < num_class; k++)
                        {
                            if (pdata[k + 5] > max_class_socre)
                            {
                                max_class_socre = pdata[k + 5];
                                max_ind = k;
                            }
                        }
                        max_class_socre *= box_score;
                        if (max_class_socre > this->confThreshold && ((max_ind == 2)) || (max_ind == 5) || (max_ind == 6) || (max_ind == 7))
                        {
                            float cx = (pdata[0] * 2.f - 0.5f + j) * stride;  ///cx
                            float cy = (pdata[1] * 2.f - 0.5f + i) * stride;   ///cy
                            float w = powf(pdata[2] * 2.f, 2.f) * anchor_w;   ///w
                            float h = powf(pdata[3] * 2.f, 2.f) * anchor_h;  ///h

                            float xmin = (cx - padw - 0.5 * w) * ratiow;
                            float ymin = (cy - padh - 0.5 * h) * ratioh;
                            float xmax = (cx - padw + 0.5 * w) * ratiow;
                            float ymax = (cy - padh + 0.5 * h) * ratioh;

                            generate_boxes.push_back(BoxInfoCar{ xmin, ymin, xmax, ymax, max_class_socre, max_ind });
                        }
                    }
                    row_ind++;
                    pdata += nout;
                }
            }
        }
    }

    // Perform non maximum suppression to eliminate redundant overlapping boxes with
    // lower confidences
    nms(generate_boxes);

    for (size_t i = 0;i < generate_boxes.size();i++)
    {
        BBox_v4  resultone;
        float x = (generate_boxes[i].x2 + generate_boxes[i].x1) / 2;
        float y = (generate_boxes[i].y1 + generate_boxes[i].y2) / 2;
        float w = generate_boxes[i].x2 - generate_boxes[i].x1;
        float h = generate_boxes[i].y2 - generate_boxes[i].y1;

        resultone.x = x;
        resultone.y = y;
        resultone.w = w;
        resultone.h = h;
        resultone.classIndex = generate_boxes[i].label;
        resultone.score = generate_boxes[i].score;
        result.push_back(resultone);


    }


}