yolov8 ncnn

目录

1. yolov8 目标检测

 2. 分类

3. 参考

---

1. yolov8 目标检测

1. 目标检测按照官方代码 8.0.81直接训练，训练后会得到 .pt 文件。

from ultralytics import YOLO
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '3'
# Load a model
model = YOLO('pcb_objectdetect_yolov8s.yaml')  # build a new model from YAML
model = YOLO('yolov8s.pt')  # load a pretrained model (recommended for training)
# model = YOLO('pcb_objectdetect_yolov8s.yaml').load('yolov8s.pt')  # build from YAML and transfer weights

# Train the model
model.train(task='detect',data='pcb_objectdetect_coco128.yaml', batch=32, workers=32, epochs=300, imgsz=640, cache=True, seed=1)

先测试下是否有问题，如果没问题，就用官方自带的 export.py 进行转换，这样参照参考把 669-486 注释，换为如下的，输出 torchscript 格式：

# def forward(self, x):
    #     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
    # 
    #     box, cls = torch.cat([xi.view(shape[0], self.no, -1) for xi in x], 2).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)

    def forward(self, x):
        z = []  # inference output
        for i in range(self.nl):
            boxes = self.cv2[i](x[i]).permute(0, 2, 3, 1)
            scores = self.cv3[i](x[i]).sigmoid().permute(0, 2, 3, 1)
            feat = torch.cat((boxes, scores), -1)
            z.append(feat)
        return tuple(z)

2. 接着用 ncnn github上提供的 pnnx，按照GitHub - pnnx/pnnx: PyTorch Neural Network eXchange 说明进行把torchscript转换为 ncnn的param 和 bin 格式

3. 接着用参考的代码，进行测试

#include "layer.h"
#include "net.h"

#include "opencv2/opencv.hpp"

#include 
#include 
#include 

#define MAX_STRIDE 32 // if yolov8-p6 model modify to 64

struct Object
{
	cv::Rect_ rect;
	int label;
	float prob;
};

static float softmax(
	const float* src,
	float* dst,
	int length
)
{
	float alpha = -FLT_MAX;
	for (int c = 0; c < length; c++)
	{
		float score = src[c];
		if (score > alpha)
		{
			alpha = score;
		}
	}

	float denominator = 0;
	float dis_sum = 0;
	for (int i = 0; i < length; ++i)
	{
		dst[i] = expf(src[i] - alpha);
		denominator += dst[i];
	}
	for (int i = 0; i < length; ++i)
	{
		dst[i] /= denominator;
		dis_sum += i * dst[i];
	}
	return dis_sum;
}

static void generate_proposals(
	int stride,
	const ncnn::Mat& feat_blob,
	const float prob_threshold,
	std::vector