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 outputfor i in range(self.nl):x[i] = self.m[i](x[i])  # convbs, _, 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_scoreconst 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;  ///cxfloat cy = (pdata[1] * 2.f - 0.5f + i) * stride;   ///cyfloat w = powf(pdata[2] * 2.f, 2.f) * anchor_w;   ///wfloat h = powf(pdata[3] * 2.f, 2.f) * anchor_h;  ///hfloat 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 confidencesnms(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);}}

yolov5 pt-＞onnx-＞om yolov5模型转onnx转om模型转换相关推荐

【YOLOv5】手把手教你使用LabVIEW ONNX Runtime部署 TensorRT加速，实现YOLOv5实时物体识别（含源码）
文章目录前言一.TensorRT简介二.准备工作三.YOLOv5模型的获取 1.下载源码 2.安装模块 3.下载预训练模型 4.转换为onnx模型四.LabVIEW使用TensorRT加速Y ...
pytorch将pt模型转onnx模型
pytorch将pt模型转onnx模型一导出ONNX模型 torch.onnx.export( model, # 要导出的模型 args, # 模型的输入参数,输入参数只需满足shape正确 on ...
机器学习实践：onnx模型转为Tensorflow2的pb模型2020
安装Tensorflow,如果已安装则检查 Tensorflow版本是否大于2.3.1 安装tensorflow-addons 从源码安装最新的onnx(2020/12),参考git 将YoloV5中 ...
yolov5笔记（3）——移动端部署自己的模型（随5.0更新）
一直以来学习目标检测的最终目标就是为了移动端的部署,比方说树莓派.jetson.安卓.ios等.之前因为实在对object_detection训练出来的东西效果不满意,所以当时没继续研究移动端部署.如 ...
【ONNX】各深度学习框架的模型转ONNX
文章目录 pytorch pytorch安装 pytorch转onnx 关于pytorch模型的题外话 cntk cntk安装 cntk转onnx mxnet mxnet安装 mxnet转onnx c ...
torch模型导出onnx
记录一下torch模型转onnx,onnx的优点是速度快,并且文件小,并且可支持较多的引擎,本文简单介绍如何导出: 首先加载我们的模型,例如加载我的模型: model = ResNetSE34V2(8 ...
编译ONNX模型Compile ONNX Models
编译ONNX模型Compile ONNX Models 本文是一篇介绍如何使用Relay部署ONNX模型的说明. 首先,必须安装ONNX包. 一个快速的解决方案是安装protobuf编译器,然后 pi ...
神经网络模型模型转ONNX
近期由于业务需要,需要将训练好的模型转为ONNX格式,为此颇费了一番功夫,在此总结一下吧.. 1.ONNX是一种神经网络模型保存的中间格式,支持多种格式的模型转为ONNX,也支持使用ONNX导入多种格 ...
Python: 从PYTORCH导出模型到ONNX，并使用ONNX运行时运行它
Python: 从PYTORCH导出模型到ONNX,并使用ONNX运行时运行它本教程我们将描述如何将PyTorch中定义的模型转换为ONNX格式,然后使用ONNX运行时运行它. ONNX运行时是一个 ...
yolov3-tiny原始weights模型转onnx模型并进行推理
时隔不知道多少天,我记起来我还有部分博客没写完(偷懒),所以不能偷懒把它完成!! 这篇博客的主要内容将yolov3-tiny.weights模型转换到.onnx模型: 使用onnnxruntime- ...

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

yolov5 pt-＞onnx-＞om yolov5模型转onnx转om模型转换相关推荐

最新文章

热门文章