Registro ONNX a NCNN

        [pytorch a onnx] pytorch-onnx

        [onnx a ncnn] onnx-ncnn

        [modelo de carga ncnn] ncnn-load

 1. Dependencias de instalación de Python

pip install onnxruntime onnx opencv-python

2. Cree un modelo y entrénelo, cargue los parámetros del modelo y genere onnx

#### pytorch 转 onnx
import torch
import torch.onnx
from model.efficientnet_model import EfficientNet_FV
model = EfficientNet_FV(num_classes=4800).cuda()

path_checkpoint = "..path\\best.ckpt"
checkpoint = torch.load(path_checkpoint)
model.load_state_dict(checkpoint)  # 加载训练好的模型

model.eval()
##### 生成 onnx 文件
x = torch.randn(1, 3, 128, 128).cuda()
with torch.no_grad():
    torch.onnx.export(
        model, x,
        "EfficientNet_b0.onnx",
        opset_version=11, 
        export_params=True,
        input_names=["input"],
        output_names=["output1"], 
    )

        Entre ellos, torch.onnx.export  es la función que viene con PyTorch para convertir el modelo a formato ONNX. Los primeros tres parámetros son el modelo a convertir, cualquier conjunto de entradas al modelo y el nombre de archivo del archivo ONNX exportado. Al convertir un modelo, es fácil entender que se requieren el modelo original y el nombre del archivo de salida, pero ¿por qué necesita proporcionar un conjunto de entradas para el modelo? Esto implica el principio de conversión ONNX. Del modelo PyTorch al modelo ONNX, es esencialmente una traducción de idioma. La idea intuitiva es analizar minuciosamente el código del modelo original como un compilador, registrando todos los flujos de control. Pero como se mencionó anteriormente, generalmente solo usamos ONNX para registrar gráficos estáticos que no consideran el flujo de control. Por lo tanto, PyTorch proporciona un método de conversión de modelos llamado trace: dado un conjunto de entradas, el modelo se vuelve a ejecutar, es decir, el gráfico de cálculo correspondiente a este conjunto de entradas se registra y guarda en formato ONNX. La función de exportación utiliza el método de exportación de seguimiento, que debe proporcionar cualquier conjunto de entradas para ejecutar el modelo. Mi imagen de prueba es de tres canales, 3*128*128 de tamaño, y aquí también se construye un tensor aleatorio de la misma forma.

        Si el código anterior se ejecuta correctamente, se agregará un archivo de modelo ONNX "EfficientNet_b0.onnx" en el directorio. Podemos usar el siguiente script para verificar que el archivo del modelo es correcto.

import onnx 
 
onnx_model = onnx.load("EfficientNet_b0.onnx") 
try: 
    onnx.checker.check_model(onnx_model) 
except Exception: 
    print("Model incorrect") 
else: 
    print("Model correct")

        Los modelos ONNX se pueden  visualizar con Netron (una herramienta de visualización de modelos de código abierto). 

3. onnx2ncnn

        Acerca de onnx2ncnn: Este es el punto clave. Después de obtener el modelo onnx, nos despedimos del framework pytorch. Ahora necesitamos convertir el modelo en formato onnx a .bin y .param que puede cargar ncnn. Utilice las herramientas ncnn compiladas anteriormente, si ncnn se compila correctamente: ncnn/build/tools/onnx/ tiene un onnx2ncnn compilado. Simplemente ingrese el siguiente comando (el modelo es el nombre de onnx), si no se genera nada, es exitoso. Por lo general, no se recomienda empaquetar partes que no sean de red en onnx.

        Aunque .param y .bin se generan con éxito, ¡por lo general falla cuando se usa ncnn para cargar el modelo! Debido a que ncnn actualmente no admite ciertas operaciones, aquí se presenta otra herramienta: ONNX Simplifier.

GitHub - daquexian/onnx-simplifier: simplifica tu modelo onnx

pip install onnxsim
onnxsim input_onnx_model output_onnx_model

Cuatro, modelo de carga ncnn

Referencia: https://blog.csdn.net/Enchanted_ZhouH/article/details/106063552

Supongo que te gusta

Origin blog.csdn.net/weixin_44855366/article/details/130169723
Recomendado
Clasificación