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 insertar y ejecutar modelos pequeños de Open Neural Network Exchange (ONNX) dentro de los módulos WebAssembly para realizar la inferencia en banda como parte de los gráficos de flujo de datos de operaciones de Azure IoT. 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.

¿Por qué usar la interferencia ONNX en banda?

Con los gráficos de flujo de datos de Azure IoT Operations, puede insertar la inferencia de un modelo ONNX pequeño directamente en la canalización en lugar de llamar a un servicio de predicción externo. Este enfoque ofrece varias ventajas prácticas:

• 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 incurre en la inferencia de CPU local, lo que evita los recorridos de ida y vuelta de red.
• Huella compacta: dirigido a modelos compactos como los modelos de clase MobileNet. Esta funcionalidad no es para modelos de transformadores grandes, aceleración de GPU/TPU o lanzamientos frecuentes de modelos A/B.
• Optimizado para casos de uso específicos:
  • En línea con el procesamiento de flujos de varios orígenes en el que las características ya están intercaladas en el grafo
  • Alineado con la semántica de tiempo de evento, por lo que la inferencia emplea las mismas marcas de tiempo que otros operadores.
  • Se mantiene lo suficientemente pequeño para integrar junto al módulo sin superar las restricciones prácticas de tamaño y memoria de WASM.
• 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 es necesario tener un registro de modelo independiente ni un cambio de punto de conexión externo.
• Restricciones de hardware: el back-end de ONNX se ejecuta en la CPU a través de la interfaz del sistema WebAssembly (WASI). wasi-nn No hay objetivos de GPU/TPU; solo se ejecutan operadores ONNX compatibles.
• 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.

Cuándo usar la inferencia ONNX dentro de la banda

Use la inferencia en banda cuando tenga estos requisitos:

• La necesidad de baja latencia para enriquecer o clasificar mensajes en línea durante la ingesta
• Modelos pequeños y eficientes, como los modelos de visión de la clase MobileNet o modelos compactos similares.
• Inferencia que debe alinearse con el procesamiento basado en el tiempo del evento y las mismas marcas de tiempo que utilizan otros operadores
• Actualizaciones sencillas mediante el envío de una nueva versión de módulo con un modelo actualizado

Evite la inferencia en banda cuando existan estos requisitos:

• Modelos de transformadores grandes, aceleración de GPU/TPU o despliegues sofisticados de A/B
• Modelos que requieren entradas tensoriales múltiples, almacenamiento de caché de valor-clave o operadores ONNX no soportados

Prerrequisitos

Antes de comenzar, asegúrese de que tiene:

• Una implementación de Azure IoT Operations con la funcionalidad de grafos de flujo de datos.
• Acceso a un registro de contenedor como Azure Container Registry.
• Entorno de desarrollo configurado para el desarrollo de módulos WebAssembly.

Para obtener instrucciones detalladas sobre la configuración, consulte Desarrollo de módulos WebAssembly.

Patrón de arquitectura

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

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 "formato" de transformación de datos: descodifica y cambia el tamaño de las imágenes a RGB24 224×224.
• Ejemplo de "instantánea" de procesamiento de imágenes y vídeos: 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, debe configurar 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 compatibilidad con WASI-NN

Para habilitar la compatibilidad con la interfaz de red neuronal de WebAssembly, agregue la función a la definición del gráfico wasi-nn:

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

Definición de operaciones y flujo de datos

Configure las operaciones que forman la canalización de inferencia. 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 sin recompilar. Estos parámetros se pasan a los módulos WASM en la inicialización:

moduleConfigurations:
  - name: module-format/map
    parameters:
      width:
        name: width
        description: "Target width for image resize (default: 224)"
        required: false
      height:
        name: height
        description: "Target height for image resize (default: 224)"
        required: false
      pixelFormat:
        name: pixel_format
        description: "Output pixel format (rgb24, bgr24, grayscale)"
        required: false

  - name: module-snapshot/map
    parameters:
      executionTarget:
        name: execution_target
        description: "Inference execution target (cpu, auto)"
        required: false
      labelMap:
        name: label_map
        description: "Label mapping strategy (embedded, imagenet, custom)"
        required: false
      scoreThreshold:
        name: score_threshold
        description: "Minimum confidence score to include in results (0.0-1.0)"
        required: false
      topK:
        name: top_k
        description: "Maximum number of predictions to return (default: 5)"
        required: false

El operador init puede leer estos valores a través de la interfaz de configuración del módulo. 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.

Directrices de preparación del modelo

Antes de insertar el modelo, asegúrese de que cumple los requisitos para la 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 de 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.

Inserción de flujo de trabajo

Siga estos pasos para insertar el modelo y los recursos asociados:

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. Integrar durante la compilación: utiliza mecanismos específicos del lenguaje para incluir los bytes del modelo en tus binarios. En Rust, use include_bytes! para datos binarios y include_str! para archivos de texto.
3. Inicializar gráfico WASI-NN: en la función init del operador, cree un grafo wasi-nn a partir de los bytes incrustados, especificando la codificación ONNX y el objetivo 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 "instantánea".

Estructura de proyecto 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

Ejemplos de referencias de archivos

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

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

Ejemplo de inserción mínima

En el ejemplo siguiente se muestra la inserción mínima de Rust. 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(())
}

Optimización del rendimiento

Para evitar volver a crear el gráfico ONNX y el contexto de ejecución de cada mensaje, inicialícelo una vez y reutilícelo. En el ejemplo público se usa un elemento estático LazyLock:

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()?;
  }
}

Implementación de los módulos

Vuelva a usar los generadores de ejemplo optimizados o compile localmente:

• Para usar el generador optimizado de Docker, consulte ejemplo de Generador de Docker de Rust.
• Para ver los pasos generales de implementación y registro de WebAssembly, consulte Uso de WebAssembly con gráficos de flujo de datos.

Siga este proceso de implementación:

1. Compila el módulo WASM en modo de liberación y genera un archivo <module-name>-<version>.wasm.
2. Envía el módulo y, opcionalmente, una definición de gráfico a su registro utilizando OCI Registry como almacenamiento (ORAS).
3. Cree o reutilice un punto de conexión del Registro en Azure IoT Operations.
4. Cree un recurso de gráfico de flujo de datos que haga referencia al artefacto de definición del gráfico.

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 de "formato" proporciona funcionalidad de descodificación y cambio de tamaño de imágenes, y el ejemplo de "instantánea" proporciona inferencia onnx y procesamiento softmax.

Para implementar este ejemplo, extraiga los artefactos del registro público, insértelos en el registro e implemente un gráfico de flujo de datos como se muestra en el ejemplo 2: Implementación de un gráfico complejo. El gráfico complejo usa estos módulos para procesar instantáneas de imagen y emitir resultados de clasificación.

Limitaciones

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

• Exclusivo para ONNX. Los gráficos de flujo de datos no admiten otros formatos como TFLite.
• Solo CPU. Sin aceleración de GPU/TPU.
• Modelos pequeños recomendados. No se admiten modelos de gran tamaño ni inferencia con uso intensivo de memoria.
• 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.
• Asegúrese de que el back-end de ONNX en el entorno de ejecución de WASM admite los operadores del modelo. Si no se admite un operador, la inferencia produce un error en el tiempo de carga o ejecución.

Pasos siguientes

• Desarrollo de módulos webAssembly
• Configuración de definiciones de grafos de WebAssembly
• Uso de WebAssembly con gráficos de flujo de datos