yolov8是yolo的最新版本,可做图像分类,目标检测,实例分割,姿态估计。
主页地址
这里测试一个分割模型。
模型如下
选yolov8n-seg模型,转成onnx,再转ncnn测试。
yolov8s-seg的ncnn版可以直接用这个
如果用python版的onnx,可以直接用notebook转,然后下载。
python版onnx代码参考
但是用notebook版本的onnx转ncnn之后,测试图片时会报错。
可以看一下转出来的yolov8n-seg.param,里面有很多的MemoryData,不够clean.
所以需要修改一些地方。(参考资料)
查了一些资料都说是要修改modules.py
.
然而本地并没有modules.py
,而是找到了modules文件夹
那么只需要在这些文件里面找到3个class并做相同的修改就可以了。
于是需要修改block.py
和head.py
两个文件。
要修改3个forward函数。
block.py
修改:
class C2f(nn.Module):
...
def forward(self, x):
"""Forward pass through C2f layer."""
#y = list(self.cv1(x).chunk(2, 1))
#y.extend(m(y[-1]) for m in self.m)
#return self.cv2(torch.cat(y, 1))
x = self.cv1(x)
x = [x, x[:, self.c:, ...]] #onnx不支持chunk操作?
x.extend(m(x[-1]) for m in self.m)
x.pop(1)
return self.cv2(torch.cat(x, 1))
head.py
中修改了Detect和Segment class.
class Detect(nn.Module):
...
def forward(self, x):
"""Concatenates and returns predicted bounding boxes and class probabilities."""
shape = x[0].shape # BCHW
for i in range(self.nl):
x[i] = torch.cat((self.cv2[i](x[i]), self.cv3[i](x[i])), 1)
if self.training:
return x
elif self.dynamic or self.shape != shape:
self.anchors, self.strides = (x.transpose(0, 1) for x in make_anchors(x, self.stride, 0.5))
self.shape = shape
#x_cat = torch.cat([xi.view(shape[0], self.no, -1) for xi in x], 2)
#if self.export and self.format in ('saved_model', 'pb', 'tflite', 'edgetpu', 'tfjs'): # avoid TF FlexSplitV ops
# box = x_cat[:, :self.reg_max * 4]
# cls = x_cat[:, self.reg_max * 4:]
#else:
# box, cls = x_cat.split((self.reg_max * 4, self.nc), 1)
#dbox = dist2bbox(self.dfl(box), self.anchors.unsqueeze(0), xywh=True, dim=1) * self.strides
#y = torch.cat((dbox, cls.sigmoid()), 1)
#return y if self.export else (y, x)
pred = torch.cat([xi.view(shape[0], self.no, -1) for xi in x], 2) #精简为这2句
return pred
class Segment(Detect):
...
def forward(self, x):
"""Return model outputs and mask coefficients if training, otherwise return outputs and mask coefficients."""
p = self.proto(x[0]) # mask protos
bs = p.shape[0] # batch size
mc = torch.cat([self.cv4[i](x[i]).view(bs, self.nm, -1) for i in range(self.nl)], 2) # mask coefficients
x = self.detect(self, x)
if self.training:
return x, mc, p
#return (torch.cat([x, mc], 1), p) if self.export else (torch.cat([x[0], mc], 1), (x[1], mc, p))
return (torch.cat([x, mc], 1).permute(0, 2, 1), p.view(bs, self.nm, -1)) if self.export else (
torch.cat([x[0], mc], 1), (x[1], mc, p))
改好了之后用如下代码export onnx文件, onnx文件会导出到和pt文件同一文件夹。
from ultralytics import YOLO
# load model
model = YOLO("pt路径/yolov8n-seg.pt")
# Export model
success = model.export(format="onnx", opset=12, simplify=True)
转好onnx之后,再把onnx转成ncnn,用onnx2ncnn命令(在build ncnn时会有这个命令)
onnx2ncnn一般在ncnn/build/tools/onnx文件夹下。
$ ./onnx2ncnn onnx路径/yolov8n-seg.onnx 想保存的路径/yolov8n-seg.param 想保存的路径/yolov8n-seg.bin
对比一下修改前和修改后导出的参数yolov8n-seg.param。
可以看到clean了很多,修改前有MemoryData, 而修改之后没有了。
下面来测试一下ncnn模型,
可以用这个代码,
需要修改几个地方:
ncnn::Mat out;
//ex.extract("output", out);
ex.extract("output0", out);
ncnn::Mat mask_proto;
//ex.extract("seg", mask_proto);
ex.extract("output1", mask_proto);
再改load_param
和load_model
的路径。
然后就可以用啦(图片来自COCO数据集)
下面是直接用yolov8n-seg.pt预测出的结果,
可以看到修改了之后效果是差了一点的(其中一个dog的框没预测出来)。文章来源:https://www.toymoban.com/news/detail-447333.html
文章来源地址https://www.toymoban.com/news/detail-447333.html
到了这里,关于yolov8seg模型转onnx转ncnn的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!