yolov8 瑞芯微RKNN和地平线Horizon芯片仿真测试部署-2023年11月15日版本

  特别说明:参考官方开源的yolov8代码、瑞芯微官方文档、地平线的官方文档,如有侵权告知删,谢谢。

  由于yolov8的官方代码结构进行了很大的调整,之前yolov8刚出来的时候写的部署博客,有网友反馈找不到对应的地方,基于截至2023年11月官方最新代码结构进行部署博客来了。写博文时使用的训练代码-版本到2023年11月15日。

  之前给出过完整的部署仿真代码和模型,今天只对如何导出能上板端芯片,且效率高的流程进行梳理。

  模型和完整仿真测试代码,放在github上参考链接 模型和代码。

  本篇博客中给出的示例,是基于船舶检测一个数据集进行训练的,只检测船舶一个类别,只是来验证流程,给出的完整代码和模型参考链接还是之前一篇博客【yolov8 瑞芯微RKNN和地平线Horizon芯片仿真测试部署】中的代码。

1 模型和训练

  训练代码参考官方开源的yolov8训练代码,如果板端芯片不支持SiLU,建议训练的时候把SiLU改为ReLU。

2 导出 yolov8 onnx

   后处理中有些算在板端芯片上效率低或者不支持,导出 onnx 需要将板端芯片不友好或不支持算子规避掉。导出onnx修改的部分。(注意以下步骤顺序不能乱,且有些步骤会运行保存,但只要能生成对应的文件就可以,报错不用管。)

第一步:
进行预测将pt只保存权重,增加代码如下图。

# 保存权重值
import torch
self.model.fuse()
self.model.eval()
print('save yolov8_relu_dict')
torch.save(self.model.state_dict(), './weights/yolov8_relu_dict.pt')

修改后运行以下代码:

from ultralytics import YOLO
model = YOLO('./weights/yolov8_relu.pt')
results = model(task='detect', mode='predict', source='./images/test.jpg', line_thickness=3, show=True, save=True, device='cpu')

第二步:
导出onnx,去除不需要的算子,修改代码如下。

class Detect(nn.Module):
    """YOLOv8 Detect head for detection models."""
    dynamic = False  # force grid reconstruction
    export = False  # export mode
    shape = None
    anchors = torch.empty(0)  # init
    strides = torch.empty(0)  # init

    def __init__(self, nc=80, ch=()):
        """Initializes the YOLOv8 detection layer with specified number of classes and channels."""
        super().__init__()
        self.nc = nc  # number of classes
        self.nl = len(ch)  # number of detection layers
        self.reg_max = 16  # DFL channels (ch[0] // 16 to scale 4/8/12/16/20 for n/s/m/l/x)
        self.no = nc + self.reg_max * 4  # number of outputs per anchor
        self.stride = torch.zeros(self.nl)  # strides computed during build
        c2, c3 = max((16, ch[0] // 4, self.reg_max * 4)), max(ch[0], min(self.nc, 100))  # channels
        self.cv2 = nn.ModuleList(
            nn.Sequential(Conv(x, c2, 3), Conv(c2, c2, 3), nn.Conv2d(c2, 4 * self.reg_max, 1)) for x in ch)
        self.cv3 = nn.ModuleList(nn.Sequential(Conv(x, c3, 3), Conv(c3, c3, 3), nn.Conv2d(c3, self.nc, 1)) for x in ch)
        self.dfl = DFL(self.reg_max) if self.reg_max > 1 else nn.Identity()

        # 导出 onnx 增加
        self.conv1x1 = nn.Conv2d(16, 1, 1, bias=False).requires_grad_(False)
        xx = torch.arange(16, dtype=torch.float)
        self.conv1x1.weight.data[:] = nn.Parameter(xx.view(1, 16, 1, 1))

    def forward(self, x):
        """Concatenates and returns predicted bounding boxes and class probabilities."""
        shape = x[0].shape  # BCHW

        # 导出 onnx 增加
        y = []
        for i in range(self.nl):
            t1 = self.cv2[i](x[i])
            t2 = self.cv3[i](x[i])
            y.append(self.conv1x1(t1.view(t1.shape[0], 4, 16, -1).transpose(2, 1).softmax(1)))
            y.append(t2)
        return y

        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

        if self.export and self.format in ('tflite', 'edgetpu'):
            # Normalize xywh with image size to mitigate quantization error of TFLite integer models as done in YOLOv5:
            # https://github.com/ultralytics/yolov5/blob/0c8de3fca4a702f8ff5c435e67f378d1fce70243/models/tf.py#L307-L309
            # See this PR for details: https://github.com/ultralytics/ultralytics/pull/1695
            img_h = shape[2] * self.stride[0]
            img_w = shape[3] * self.stride[0]
            img_size = torch.tensor([img_w, img_h, img_w, img_h], device=dbox.device).reshape(1, 4, 1)
            dbox /= img_size

        y = torch.cat((dbox, cls.sigmoid()), 1)
        return y if self.export else (y, x)

    def bias_init(self):
        """Initialize Detect() biases, WARNING: requires stride availability."""
        m = self  # self.model[-1]  # Detect() module
        # cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1
        # ncf = math.log(0.6 / (m.nc - 0.999999)) if cf is None else torch.log(cf / cf.sum())  # nominal class frequency
        for a, b, s in zip(m.cv2, m.cv3, m.stride):  # from
            a[-1].bias.data[:] = 1.0  # box
            b[-1].bias.data[:m.nc] = math.log(5 / m.nc / (640 / s) ** 2)  # cls (.01 objects, 80 classes, 640 img)

增加保存onnx模型代码,如下:

    def _new(self, cfg: str, task=None, model=None, verbose=True):
        """
        Initializes a new model and infers the task type from the model definitions.
        """
        cfg_dict = yaml_model_load(cfg)
        self.cfg = cfg
        self.task = task or guess_model_task(cfg_dict)
        self.model = (model or self._smart_load('model'))(cfg_dict, verbose=verbose and RANK == -1)  # build model
        self.overrides['model'] = self.cfg
        self.overrides['task'] = self.task

        # Below added to allow export from YAMLs
        self.model.args = {**DEFAULT_CFG_DICT, **self.overrides}  # combine default and model args (prefer model args)
        self.model.task = self.task

        # 保存onnx
        import torch
        self.model.fuse()
        self.model.load_state_dict(torch.load('./weights/yolov8_relu_dict.pt', map_location='cpu'), strict=False)
        self.model.eval()

        print("===========  onnx =========== ")
        dummy_input = torch.randn(1, 3, 640, 640)
        input_names = ["data"]
        output_names = ["reg1", "cls1", "reg2", "cls2", "reg3", "cls3"]
        torch.onnx.export(self.model, dummy_input, "./weights/yolov8_relu.onnx", verbose=False, input_names=input_names, output_names=output_names, opset_version=12)
        print("======================== convert onnx Finished! .... ")

以上修改后完运行以下代码(注意和第一次运行的不一样,这次加载的是yaml):

from ultralytics import YOLO

model = YOLO('./ultralytics/models/v8/yolov8n.yaml')

  特别说明顺序一定不能错,一定不能错,如果重新训练模型后,需要再次生成onnx模型,需要先把第二步的修改撤回再运行第一步。

3 pytorch 测试效果

3 onnx 测试效果


注:图片来源一求助网友,若侵权请告知删除。

4 rknn 板端C++部署

C++完整部署参考

说明:由于检测的类别数目不一样,不同的检测类别数目,模型推理和后处理的时耗也是不完全一致的,类别越少模型推理和后处理时耗越少。

5 yolov8导出瑞芯微rknn和地平线horizon仿真测

5.1 瑞芯微 rknn 仿真

  瑞芯微环境搭建和详细步骤参考上一篇 【瑞芯微RKNN模型转换和PC端仿真】。
  yolov8导出rknn模型代码和后处理参考 yolov8_rknn

5.2 地平线仿真

  地平线环境搭建和详细步骤参考上一篇 【地平线Horizon模型转换和PC端仿真测试】。
  yolov8导出地平线模型代码和后处理参考 yolov8_horizon

6 官方导出onnx方式进行瑞芯微rknn和地平线horizon仿真测试

yolov8 官方模型进行瑞芯微RKNN和地平线Horizon芯片仿真测试部署

文章出处登录后可见!

已经登录?立即刷新

共计人评分,平均

到目前为止还没有投票!成为第一位评论此文章。

(0)
扎眼的阳光的头像扎眼的阳光普通用户
上一篇 2023年12月19日
下一篇 2023年12月19日

相关推荐