Executar inferência ONNX em grafos de fluxo de dados WASM (WebAssembly)

Este artigo mostra como inserir e executar pequenos modelos do ONNX (Open Neural Network Exchange) em seus módulos WebAssembly para executar a inferência em banda como parte dos grafos de fluxo de dados das Operações de IoT do Azure. Use essa abordagem para enriquecimento de baixa latência e classificação diretamente em dados de streaming sem chamar serviços de previsão externa.

Por que usar a interferência ONNX na banda

Com grafos de fluxo de dados de operações do Azure IoT, você pode inserir a inferência de um pequeno modelo ONNX diretamente no pipeline, em vez de chamar um serviço de previsão externo. Essa abordagem oferece várias vantagens práticas:

• Baixa latência: execute o enriquecimento ou a classificação em tempo real na mesma via do operador por onde os dados chegam. Cada mensagem incorre apenas na inferência da CPU local, evitando viagens de ida e volta de rede.
• Volume compacto: modelos compactos de destino, como modelos da classe MobileNet. Essa funcionalidade não é para modelos de transformer grandes, aceleração de GPU/TPU ou implementações frequentes de modelos A/B.
• Otimizado para casos de uso específicos:
  • Em linha com o processamento de fluxo de várias fontes, onde os recursos já estão colocados no grafo.
  • Alinhado com a semântica do tempo de evento para que a inferência use os mesmos timestamps que outros operadores.
  • Mantido pequeno o suficiente para ser incorporado ao módulo sem exceder os limites práticos de tamanho e memória do WASM
• Atualizações simples: envie um novo módulo com WASM e modelo inserido e atualize a referência de definição de grafo. Não é necessário ter um registro de modelo separado ou uma alteração de ponto de extremidade externo.
• Restrições de hardware: o back-end ONNX é executado na CPU por meio da WASI (Interface do Sistema WebAssembly). wasi-nn Nenhum destino de GPU/TPU; somente operadores ONNX com suporte são executados.
• Dimensionamento horizontal: a inferência é dimensionada conforme o grafo de fluxo de dados é dimensionado. Quando o runtime adiciona mais trabalhadores para aumentar a taxa de transferência, cada trabalhador carrega o modelo incorporado e participa do balanceamento de carga.

Quando usar a inferência ONNX em banda

Use a inferência em banda quando você tiver os seguintes requisitos:

• A baixa latência precisa enriquecer ou classificar mensagens em linha no momento de ingestão
• Modelos pequenos e eficientes, como os modelos de visão da classe MobileNet ou modelos compactos semelhantes
• A inferência que precisa se alinhar com o processamento em tempo de evento e os mesmos carimbos de data/hora que outros operadores
• Atualizações simples enviando uma nova versão de módulo com um modelo atualizado

Evite a inferência em banda quando tiver estes requisitos:

• Modelos de transformadores de grande porte, aceleração de GPU/TPU ou lançamentos A/B sofisticados
• Os modelos que exigem várias entradas de tensores, cache de chave-valor ou operadores ONNX sem suporte

Pré-requisitos

Antes de começar, verifique se você tem:

• Uma implantação de Operações de IoT do Azure com funcionalidade de grafos de fluxo de dados.
• Acesso a um registro de contêiner como o Registro de Contêiner do Azure.
• Ambiente de desenvolvimento configurado para o desenvolvimento do módulo WebAssembly.

Para obter instruções detalhadas de instalação, consulte Desenvolver módulos WebAssembly.

Padrão de arquitetura

O padrão comum para inferência ONNX em grafos de fluxo de dados inclui:

1. Dados de pré-processamento: transforme dados de entrada brutos para corresponder ao formato esperado do modelo. Para modelos de imagem, esse processo normalmente envolve:
   • Decodificação de bytes de imagem.
   • Redimensionando para uma dimensão de destino (por exemplo, 224×224).
   • Convertendo o espaço de cores (por exemplo, RGB em BGR).
   • Normalizando valores de pixel para o intervalo esperado (0 a 1 ou -1 a 1).
   • Organizando dados no layout de tensor correto: NCHW (lote, canais, altura, largura) ou NHWC (lote, altura, largura, canais).
2. Executar inferência: converta dados pré-processados em tensores usando a interface wasi-nn, carregue o modelo ONNX incorporado com o back-end da CPU, defina tensores de entrada no contexto de execução, invoque o processamento para frente do modelo e recupere tensores de saída que contêm previsões brutas.
3. Saídas de pós-processamento: transformar saídas de modelo bruto em resultados significativos. Operações comuns:
   • Aplique softmax para produzir probabilidades de classificação.
   • Selecione as previsões top-K.
   • Aplique um limite de confiança para filtrar os resultados de baixa confiança.
   • Mapear os índices de previsão para rótulos compreensíveis por humanos.
   • Formatar resultados para consumo posterior.

Nos exemplos de IoT para operadores RUST WASM , você pode encontrar dois exemplos que seguem este padrão:

• Amostra de "formato" da transformação de dados: decodifica e redimensiona imagens para RGB24 224×224.
• Exemplo de "instantâneo" de processamento de imagem/vídeo: ele incorpora um modelo MobileNet v2 ONNX, executa inferência na CPU e calcula softmax.

Configurar a definição do grafo

Para habilitar a inferência ONNX no grafo de fluxo de dados, você precisa configurar a estrutura do grafo e os parâmetros do módulo. A definição do grafo especifica o fluxo de pipeline, enquanto as configurações dos módulos permitem a personalização em tempo de execução do comportamento de pré-processamento e inferência.

Habilitar o suporte WASI-NN

Para habilitar o suporte à interface WebAssembly Neural Network, adicione o recurso wasi-nn à definição do grafo.

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

Definir operações e fluxo de dados

Configure as operações que formam o pipeline de inferência. Este exemplo mostra um fluxo de trabalho de classificação de imagem típico:

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" }

Essa configuração cria um pipeline em que:

• camera-input recebe dados brutos de imagem de uma fonte
• module-format/map pré-processa imagens (decodificar, redimensionar, formatar conversão)
• module-snapshot/map executa a inferência e o pós-processamento do ONNX
• results-output emite resultados de classificação para um coletor

Configurar parâmetros de módulo

Defina parâmetros de runtime para personalizar o comportamento do módulo sem recompilar. Esses parâmetros são passados para os módulos WASM na inicialização:

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

Seu operador init pode ler esses valores por meio da interface de configuração do módulo. Para obter detalhes, consulte os parâmetros de configuração do módulo.

Empacotar o modelo

A incorporação de modelos ONNX diretamente em seu componente WASM garante a implantação atômica e a consistência da versão. Essa abordagem simplifica a distribuição e remove as dependências de runtime em arquivos ou registros de modelo externos.

Diretrizes de preparação do modelo

Antes de inserir seu modelo, verifique se ele atende aos requisitos de implantação do WASM:

• Mantenha modelos abaixo de 50 MB para tempos práticos de carregamento de WASM e restrições de memória.
• Verifique se o modelo aceita uma única entrada tensor em um formato comum (float32 ou uint8).
• Verifique se o back-end de runtime WASM ONNX dá suporte a todos os operadores usados pelo modelo.
• Use ferramentas de otimização ONNX para reduzir o tamanho do modelo e melhorar a velocidade de inferência.

Inserindo fluxo de trabalho

Siga estas etapas para inserir seu modelo e recursos associados:

1. Organizar ativos de modelo: coloque o arquivo de modelo .onnx e os labels.txt opcionais em sua árvore de origem. Use uma estrutura de diretório dedicada, como src/fixture/models/ e src/fixture/labels/ para uma organização clara.
2. Incorporar em tempo de compilação: Use mecanismos específicos da linguagem para incluir os bytes do modelo em seu binário. No Rust, use include_bytes! para dados binários e include_str! para arquivos de texto.
3. Inicializar grafo WASI-NN: na função init do operador, crie um grafo wasi-nn a partir dos bytes incorporados, especificando a codificação ONNX e o destino de execução para CPU.
4. Implementar o loop de inferência: para cada mensagem de entrada, pré-processe entradas para corresponder aos requisitos do modelo, defina tensores de entrada, execute inferência, recupere saídas e aplique o pós-processamento.
5. Lidar com erros normalmente: implemente o tratamento de erros adequado para falhas de carregamento de modelo, operadores sem suporte e erros de inferência de runtime.

Para obter um padrão de implementação completo, consulte o exemplo de "instantâneo".

Estrutura de projeto recomendada

Organize seu projeto de módulo WASM com clara separação de preocupações:

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

Exemplos de referências de arquivo

Use o seguinte layout de arquivo do exemplo de "instantâneo" como referência:

• Diretório de rótulos – Contém vários arquivos de mapeamento de rótulo
• Diretório de modelos – Contém arquivos e metadados do modelo ONNX

Exemplo de incorporação mínima

O exemplo a seguir mostra a inserção mínima do Rust. Os caminhos são relativos ao arquivo de origem que contém a 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(())
}

Otimização de desempenho

Para evitar recriar o grafo ONNX e o contexto de execução para cada mensagem, inicialize-o uma vez e reutilize-o. O exemplo público usa uma estática 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()?;
  }
}

Implante seus módulos

Reutilize os construtores de exemplo simplificados ou crie localmente:

• Para usar o construtor simplificado do Docker, consulte o exemplo do Construtor do Rust Docker
• Para obter uma implantação geral do WebAssembly e etapas do registro, consulte Usar WebAssembly com grafos de fluxo de dados

Siga este processo de implantação:

1. Crie seu módulo WASM em modo de release e produza um arquivo <module-name>-<version>.wasm.
2. Envie o módulo e, opcionalmente, uma definição de gráfico para o registro usando o Registro OCI como Armazenamento (ORAS).
3. Crie ou reutilize um ponto de extremidade de registro nas Operações de IoT do Azure.
4. Crie um recurso de grafo de fluxo de dados que faça referência ao artefato de definição de grafo.

Exemplo: classificação de imagem da MobileNet

Os exemplos públicos de IoT fornecem dois exemplos conectados a um grafo para classificação de imagem: o exemplo de "formato" fornece a funcionalidade de decodificação e redimensionamento de imagem, e o exemplo de "instantâneo" fornece inferência ONNX e processamento softmax.

Para implantar esse exemplo, efetue pull dos artefatos do registro público, efetue push para o seu registro e implante um grafo de fluxo de dados conforme mostrado no Exemplo 2: Implantar um grafo complexo. O grafo complexo usa esses módulos para processar instantâneos de imagem e emitir resultados de classificação.

Limitações

A inferência em grafos de fluxo de dados WASM tem as seguintes limitações:

• Apenas ONNX. Os grafos de fluxo de dados não dão suporte a outros formatos, como o TFLite.
• Somente CPU. Nenhuma aceleração de GPU/TPU.
• Modelos pequenos recomendados. Não há suporte para modelos grandes e inferência com uso intensivo de memória.
• Há suporte para modelos de entrada de tensor único. Não há suporte para modelos com várias entradas, armazenamento em cache de chave-valor e cenários avançados de sequência ou gerativos.
• Verifique se o backend ONNX no tempo de execução do WASM dá suporte aos operadores do modelo. Se não houver suporte para um operador, a inferência falhará no tempo de carregamento ou execução.

Próximas etapas

• Desenvolver módulos WebAssembly
• Configurar definições de grafo do WebAssembly
• Usar WebAssembly com grafos de fluxo de dados