Ejecución de la inferencia de ONNX en gráficos de flujo de datos de WebAssembly (WASM)

En este artículo se muestra cómo incrustar y ejecutar pequeños modelos Open Neural Network Exchange (ONNX) dentro de sus módulos de WebAssembly para realizar inferencia en banda como parte de los gráficos de flujo de datos de Operaciones de IoT de Azure. Use este enfoque para el enriquecimiento y la clasificación de baja latencia directamente en los datos de streaming sin llamar a servicios de predicción externos.

Prerrequisitos

Antes de empezar, asegúrese de que tiene los siguientes elementos:

• Una implementación de Operaciones de IoT de Azure con la funcionalidad de gráficos de flujo de datos.
• Un registro de contenedor como Azure Container Registry y un punto de conexión del registro de contenedor configurado. Para más información, consulte Configuración de puntos de conexión del Registro.
• Un entorno de desarrollo configurado para el desarrollo de módulos WebAssembly. Para obtener opciones e instrucciones detalladas de configuración del entorno, consulte Desarrollo de módulos WebAssembly.

Ventajas de la inferencia ONNX en banda

Con los gráficos de flujo de datos de Operaciones de IoT de Azure, puede insertar una pequeña inferencia de modelos ONNX directamente en el flujo de datos en lugar de llamar a un servicio de predicción externo. Este enfoque ofrece las siguientes ventajas:

• Baja latencia: realice el enriquecimiento en tiempo real o la clasificación en la misma ruta de acceso del operador donde llegan los datos. Cada mensaje solo requiere inferencia local en la CPU, lo que evita los viajes de ida y vuelta por la red.
• Integrado con el procesamiento de flujos: ejecute la inferencia junto con el procesamiento de flujos de múltiples fuentes, donde las variables ya se encuentran ubicadas conjuntamente en el grafo, y alinéela con la semántica de tiempo de evento para que la inferencia use las mismas marcas de tiempo que otros operadores.
• Actualizaciones sencillas: envíe un nuevo módulo con WASM y modelo incrustado y, a continuación, actualice la referencia de definición del grafo. No necesita un registro de modelo independiente ni un cambio de punto de conexión externo.
• Escalado horizontal: la inferencia se escala a medida que se escala el gráfico de flujo de datos. Cuando el entorno de tiempo de ejecución agrega más trabajadores para el rendimiento de procesamiento, cada trabajador carga el modelo incrustado y participa en el balanceo de carga.

La inferencia se ejecuta en una CPU a través de la interfaz del sistema WebAssembly (WASI) wasi-nn . Para conocer los formatos de modelo admitidos y las restricciones de hardware, consulte Limitaciones de la inferencia de ONNX en gráficos de flujo de datos WASM.

¿Cuándo debe usar la inferencia ONNX en banda?

Use la inferencia ONNX en banda en gráficos de flujo de datos cuando necesite las siguientes funcionalidades:

• Baja latencia para enriquecer o clasificar mensajes en línea durante la ingesta.
• Modelos de visión pequeños y eficientes, como los de tipo MobileNet.
• Inferencia que se alinea con procesamiento de hora del evento y usa las mismas marcas de tiempo que otros operadores.
• Actualizaciones de modelos simples mediante el envío de una nueva versión de módulo.

No use la inferencia ONNX en banda cuando necesite las siguientes funcionalidades:

• Modelos de transformadores grandes, aceleración de GPU o TPU o implementaciones sofisticadas de A/B.
• Modelos que requieren varias entradas de tensor, almacenamiento en caché de clave-valor o operadores ONNX no admitidos.

Patrón de arquitectura para la inferencia de ONNX en gráficos de flujo de datos

El patrón común para la inferencia de ONNX en gráficos de flujo de datos incluye las siguientes fases:

1. Preprocesar datos: transforme los datos de entrada sin procesar para que coincidan con el formato esperado del modelo. En el caso de los modelos de imagen, este proceso normalmente implica:
   • Descodificación de bytes de imagen.
   • Redimensionar a una dimensión objetivo (por ejemplo, 224×224).
   • Convertir el espacio de color (por ejemplo, RGB a BGR).
   • Normalizar valores de píxeles en el intervalo esperado (de 0 a 1 o -1 a 1).
   • Organizar los datos en el diseño correcto de tensor: NCHW (lote, canales, altura, ancho) o NHWC (lote, altura, ancho, canales).
2. Ejecutar inferencia: convierta los datos preprocesados en tensores mediante la interfaz wasi-nn, cargue el modelo ONNX incrustado con el backend de CPU, establezca tensores de entrada en el contexto de ejecución, invoque el paso forward del modelo y recupere tensores de salida que contienen predicciones sin procesar.
3. Salidas posteriores al procesamiento: transforme las salidas del modelo sin procesar en resultados significativos. Operaciones comunes:
   • Aplique softmax para generar probabilidades de clasificación.
   • Seleccione predicciones de nivel superior K.
   • Aplique un umbral de confianza para filtrar los resultados de confianza baja.
   • Asigne índices de predicción a etiquetas legibles por humanos.
   • Dar formato a los resultados para su uso posterior.

En los ejemplos de IoT para operadores WASM de Rust puede encontrar dos ejemplos que siguen este patrón:

• Ejemplo de transformación "formato" de datos: descodifica y cambia el tamaño de las imágenes a RGB24 224×224.
• Ejemplo de "instantánea" de procesamiento de imagen/vídeo: inserta un modelo de ONNX de MobileNet v2, ejecuta la inferencia de CPU y calcula softmax.

Configuración de la definición del grafo

Para habilitar la inferencia de ONNX en el gráfico de flujo de datos, configure tanto la estructura del grafo como los parámetros del módulo. La definición del grafo especifica el flujo de canalización, mientras que las configuraciones de módulo permiten la personalización en tiempo de ejecución del comportamiento de preprocesamiento e inferencia.

Habilitación de la característica wasi-nn

Para habilitar la interfaz de red neuronal de WebAssembly (wasi-nn) para la inferencia con ONNX, añada la característica wasi-nn a la definición del grafo:

moduleRequirements:
  apiVersion: "1.1.0"
  runtimeVersion: "1.1.0"
  features:
    - name: "wasi-nn"

Definir operaciones para el canal de inferencia

Configure las operaciones que forman la canalización de inferencia de ONNX. En este ejemplo se muestra un flujo de trabajo típico de clasificación de imágenes:

operations:
  - operationType: "source"
    name: "camera-input"
  - operationType: "map"
    name: "module-format/map"
    module: "format:1.0.0"
  - operationType: "map"
    name: "module-snapshot/map"
    module: "snapshot:1.0.0"
  - operationType: "sink"
    name: "results-output"

connections:
  - from: { name: "camera-input" }
    to: { name: "module-format/map" }
  - from: { name: "module-format/map" }
    to: { name: "module-snapshot/map" }
  - from: { name: "module-snapshot/map" }
    to: { name: "results-output" }

Esta configuración crea una canalización donde:

• camera-input recibe datos de imagen sin procesar de un origen
• module-format/map preprocesa imágenes (descodificación, cambio de tamaño y conversión de formato)
• module-snapshot/map ejecuta la inferencia y el postprocesamiento de ONNX
• results-output emite resultados de clasificación a un receptor

Configuración de parámetros de módulo

Defina parámetros en tiempo de ejecución para personalizar el comportamiento del módulo WASM sin necesidad de recompilar. Estos parámetros pasan a los módulos WASM en la inicialización.

Para obtener más información, consulte Parámetros de configuración del módulo.

Empaquetar el modelo

La inserción de modelos ONNX directamente en el componente WASM garantiza la implementación atómica y la coherencia de la versión. Este enfoque simplifica la distribución y quita las dependencias en tiempo de ejecución en archivos o registros de modelos externos.

Requisitos de preparación del modelo ONNX

Antes de insertar el modelo, asegúrese de que cumple los requisitos de implementación de WASM:

• Mantenga los modelos menores de 50 MB para las restricciones prácticas de carga y memoria de WASM.
• Compruebe que el modelo acepta una sola entrada tensor en un formato común (float32 o uint8).
• Compruebe que el back-end en tiempo de ejecución de ONNX de WASM admite todos los operadores que usa el modelo.
• Use las herramientas de optimización de ONNX para reducir el tamaño del modelo y mejorar la velocidad de inferencia.

Pasos para insertar un modelo ONNX en un módulo WASM

Siga estos pasos para insertar el modelo ONNX y los recursos asociados en un módulo WASM:

1. Organizar los recursos del modelo: coloque el archivo modelo y los archivos opcionales .onnx en el árbol de origen. Use una estructura de directorios dedicada como src/fixture/models/ y src/fixture/labels/ para una organización clara.
2. Incrustar en tiempo de compilación: Use herramientas específicas del lenguaje para incluir los bytes del modelo en el binario. En Rust, use include_bytes! para datos binarios y include_str! para archivos de texto.
3. Inicializar el grafo wasi-nn: En la función init de su operador, cree un grafo wasi-nn a partir de los bytes incrustados, especificando la codificación ONNX y el destino de ejecución en CPU.
4. Implementar bucle de inferencia: Para cada mensaje entrante, preprocesar las entradas para que coincidan con los requisitos del modelo, establecer los tensores de entrada, ejecutar la inferencia, recuperar las salidas y aplicar postprocesamiento.
5. Gestione errores de manera organizada: implemente el manejo de errores adecuado para fallas de carga del modelo, operadores no admitidos y errores de inferencia durante la ejecución.

Para obtener un patrón de implementación completo, consulte el ejemplo "snapshot".

Estructura de proyecto WASM recomendada

Organice el proyecto del módulo WASM con una separación clara de los problemas:

src/
├── lib.rs                 # Main module implementation
├── model/
│   ├── mod.rs            # Model management module
│   └── inference.rs      # Inference logic
└── fixture/
    ├── models/
    │   ├── mobilenet.onnx      # Primary model
    │   └── mobilenet_opt.onnx  # Optimized variant
    └── labels/
        ├── imagenet.txt        # ImageNet class labels
        └── custom.txt          # Custom label mappings

Diseño de archivos de ejemplo de la instantánea de ejemplo

Use el siguiente diseño de archivo del ejemplo de "instantánea" como referencia:

• directorio Labels: contiene varios archivos de asignación de etiquetas
• directorio Models: contiene metadatos y archivos de modelo ONNX

Ejemplo mínimo de Rust para insertar un modelo ONNX

En el siguiente ejemplo de Rust se muestra el código mínimo necesario para insertar un modelo ONNX en un módulo WASM. Las rutas de acceso son relativas al archivo de origen que contiene la macro:

// src/lib.rs (example)
// Embed ONNX model and label map into the component
static MODEL: &[u8] = include_bytes!("fixture/models/mobilenet.onnx");
static LABEL_MAP: &[u8] = include_bytes!("fixture/labels/synset.txt");

fn init_model() -> Result<(), anyhow::Error> {
  // Create wasi-nn graph from embedded ONNX bytes using the CPU backend
  // Pseudocode – refer to the snapshot sample for the full implementation
  // use wasi_nn::{graph::{load, GraphEncoding, ExecutionTarget}, Graph};
  // let graph = load(&[MODEL.to_vec()], GraphEncoding::Onnx, ExecutionTarget::Cpu)?;
  // let exec_ctx = Graph::init_execution_context(&graph)?;
  Ok(())
}

Reutilice el gráfico ONNX y el contexto de ejecución entre mensajes

Para evitar volver a crear el gráfico ONNX y el contexto de ejecución de cada mensaje, inicialícelos una vez y reutilícelos. El ejemplo de instantánea pública usa un LazyLock estático para inicializar el grafo y el contexto de ejecución una vez por trabajador:

use crate::wasi::nn::{
    graph::{load, ExecutionTarget, Graph, GraphEncoding, GraphExecutionContext},
    tensor::{Tensor, TensorData, TensorDimensions, TensorType},
};

static mut CONTEXT: LazyLock<GraphExecutionContext> = LazyLock::new(|| {
    let graph = load(&[MODEL.to_vec()], GraphEncoding::Onnx, ExecutionTarget::Cpu).unwrap();
    Graph::init_execution_context(&graph).unwrap()
});

fn run_inference(/* input tensors, etc. */) {
    unsafe {
        // (*CONTEXT).compute()?;
    }
}

Depure y pruebe su módulo ONNX localmente

Antes de implementar en Operaciones de IoT de Azure, pruebe el módulo de inferencia de ONNX localmente para validar la funcionalidad y el rendimiento. Para obtener más información, consulte:

• Depurar módulos WASM en VS Code
• Prueba de módulos WASM

Configuración e implementación del módulo ONNX

Cuando esté listo para usar el módulo de inferencia de ONNX en un flujo de datos o conector de Operaciones de IoT de Azure, siga estos pasos:

• Configuración de definiciones de grafos de WebAssembly (WASM) para gráficos y conectores de flujo de datos
• Implementación de módulos y definiciones de grafos de WebAssembly (WASM)

Ejemplo: Clasificación de imágenes de MobileNet

Los ejemplos públicos de IoT proporcionan dos ejemplos conectados a un grafo para la clasificación de imágenes:

• El ejemplo "format" proporciona la funcionalidad de descodificación y cambio de tamaño de la imagen.
• El ejemplo "snapshot" proporciona inferencia onnx y procesamiento softmax.

Para obtener más información sobre y ejecutar el ejemplo que usa estos módulos, vea Ejemplo 2: Implementación de un grafo complejo.

Limitaciones de la inferencia de ONNX en gráficos de flujo de datos WASM

La inferencia en gráficos de flujo de datos WASM tiene las siguientes limitaciones:

• Formato de modelo: solo ONNX. Los gráficos de flujo de datos no admiten otros formatos como TFLite.
• Hardware: solo CPU. No se admite la aceleración de GPU y TPU.
• Tamaño del modelo: no se admiten modelos grandes ni inferencia con uso intensivo de memoria. Use modelos pequeños, como arquitecturas de clase MobileNet.
• Entradas del modelo: solo se admiten modelos de entrada de tensor único. No se admiten modelos de entrada múltiple, almacenamiento en caché de clave-valor ni escenarios de secuencia o generativos avanzados.
• Compatibilidad con operadores: el back-end de ONNX en el entorno de ejecución de WASM debe admitir todos los operadores que usa el modelo. Si no se admite un operador, la inferencia produce un error en el tiempo de carga o ejecución.

Contenido relacionado

• Desarrollo de módulos webAssembly
• Configuración de definiciones de grafos de WebAssembly (WASM) para gráficos y conectores de flujo de datos
• Implementación de módulos y definiciones de grafos de WebAssembly (WASM)