Convertir yolov5 a modelo rknn

contenido

1 método de conversión

2 problema de versión yolov5

Problema de formato de salida de 3 yolov5

4 Problemas de formato al convertir al modelo rknn


1 método de conversión

Convierta el modelo pth de yolov5 en el modelo rknn, el método específico se divide en dos pasos

  1. Use export.py que viene con el proyecto yolov5 para convertir el modelo pt al modelo onnx
  2. Use el siguiente script para generar el modelo rknn

Convierta el modelo onnx al modelo rknn de la plataforma RV1126 - cumtchw - CSDN Blog

Esta nota no es para introducir el método de conversión específico, sino para registrar algunos problemas en el proceso de conversión intermedio.

2 problema de versión yolov5

Queremos usar C++ para inferir el modelo rknn de yolov5 en la placa de Rockchip Actualmente, generalmente usamos la demostración de C++ proporcionada por Rockchip.

Demostración oficial de razonamiento YOLOv5 de Rockchip RV1126

El procesamiento posterior de esta demostración está dirigido al modelo de tres salidas. Si queremos usar esta demostración, entonces el yolov5 que usamos debe usar la versión v5.0 , no la versión v6.0, porque el modelo de la versión v5.0 tiene tres salidas, la versión v6.0 finalmente combina las tres salidas.

Problema de formato de salida de 3 yolov5

Hay tres salidas del yolov5-5.0 estándar, que son

1x255x80x80

1x255x40x40

1x255x20x20

 Entre ellos, 255 aquí es 85*3, donde 3 se refiere a 3 cajas generadas por 3 anclas ( no RGB de tres canales, y no hay un concepto RGB en la salida final ), y 85 aquí se refiere a 5+80= 85, donde 80 es el número de categorías, cada categoría corresponde a un puntaje de etiqueta, un total de 80 puntajes de etiqueta, y 5 se refiere a las cuatro coordenadas del cuadro más un puntaje de cuadro.

Si es la versión yolov5-6.0, 6.0 combina las tres salidas, entonces el arreglo es xywh box_score id Id1_score id2_score .... idn_score, luego si 3*80*80*12, el número de categorías de salida es 6 no 7, porque hay una identificación en él, que es equivalente a calcular la identificación, y las tres salidas de la versión yolov5-5.0 no se combinan, solo hay puntajes de categoría y no categorías.

4 Problemas de formato al convertir al modelo rknn

Aquí tomamos como ejemplo un modelo de reconocimiento de máscara con 3 etiquetas (sin máscara, usando una máscara correctamente, usando una máscara incorrectamente). La salida del modelo de reconocimiento de máscara con 3 etiquetas es la siguiente, de las cuales 24 son la salida de 3 canales * 8, y la salida de 8 es el xywh box_score Id1_score id2_score id3_score;

1*24*80*80

1*24*40*40

1*24*20*20

Cuando convertí por primera vez antes, descubrí que el resultado del modelo rknn transferido era completamente incorrecto al razonar en la placa RV1126 de Rockchip. Había varios cuadros pequeños en toda la imagen. Después de buscar, descubrí que mi formato de salida no era 1. *24*20*20, pero se convirtió en 1*3*20*20*8, así que modifiqué yolo.py en el proyecto yolov5-5.0 de la siguiente manera antes de usar export.py para generar el modelo onnx, y luego use export.py para generar el modelo onnx, y luego conviértalo al modelo rknn.

Tenga en cuenta que la modificación consiste en eliminar la permutación, y view() se multiplica por self.na y self.no.

O modifíquelo de la siguiente manera, comente self.training |= self.export, y luego agregue una rama de exportación directamente:

    def forward(self, x):
        # x = x.copy()  # for profiling
        z = []  # inference output
       # self.training |= self.export
        if self.export:
            for i in range(self.nl):
                x[i] = self.m[i](x[i])
                bs, _, ny, nx = x[i].shape  # x(bs,48,20,20) to x(bs,3,20,20,16)
                x[i] = x[i].view(bs, self.na*self.no, ny, nx).contiguous()

            return x

        for i in range(self.nl):
            x[i] = self.m[i](x[i])  # conv
            bs, _, 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()
        .......后面代码........

El formato de salida después de esta modificación es el siguiente:

Cuando netron mira el modelo onnx

 Cuando netron mira el modelo rknn:

 La impresión cuando se ejecuta en la placa RV1126 es la siguiente:

Supongo que te gusta

Origin blog.csdn.net/u013171226/article/details/123401225
Recomendado
Clasificación