yolov8seg模型转onnx转ncnn

news2024/10/6 10:33:46

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);