Noções Básicas de ML em Rust

10 exemplos para você começar com machine learning em Rust - 7 básicos e 3 intermediários.

Pré-requisitos

cargo new ml-playground && cd ml-playground
cargo add candle-core candle-nn --features mkl
cargo add tokenizers
cargo add ort --features download-binaries
cargo add axum tokio serde --features derive

Ferramentas: Os exemplos visam Rust 1.97.0 (edição 2024). Fixe as versões dos crates de ML em produção - as APIs evoluem rapidamente.

Exemplos Básicos

1. Carregar um modelo ONNX com ort

Execute modelos exportados sem um runtime Python no momento da inferência.

use ort::session::Session;
 
fn main() -> ort::Result<()> {
    let session = Session::builder()?.commit_from_file("model.onnx")?;
    println!("inputs: {:?}", session.inputs);
    Ok(())
}

• ONNX é o caminho de exportação comum de PyTorch, sklearn e TensorFlow.
• ort agrupa o ONNX Runtime para execução em CPU/GPU.
• Inspecione os nomes e formas das entradas antes de construir tensores.

Relacionado: ONNX Runtime (ort) - loop de inferência completo

2. Tokenizar texto para um LLM

Converta prompts em IDs de token com o crate tokenizers do Hugging Face.

use tokenizers::Tokenizer;
 
fn main() -> tokenizers::Result<()> {
    let tok = Tokenizer::from_file("tokenizer.json")?;
    let enc = tok.encode("Hello, Rust ML!", true)?;
    println!("ids: {:?}", enc.get_ids());
    Ok(())
}

• tokenizer.json vem com a maioria dos repositórios de modelos HF.
• add_special_tokens controla a inserção de BOS/EOS para modelos decodificadores.
• A tokenização é frequentemente 10-30% da latência de LLM - mantenha-a em Rust ao lado da inferência.

Relacionado: Tokenizers - batching e truncamento

3. Tensor simples com candle

Crie e multiplique tensores na CPU (ou GPU com recursos habilitados).

use candle_core::{Device, Tensor};
 
fn main() -> candle_core::Result<()> {
    let device = Device::Cpu;
    let a = Tensor::new(&[[1f32, 2.0], [3.0, 4.0]], &device)?;
    let b = Tensor::new(&[[2f32, 0.0], [1.0, 2.0]], &device)?;
    let c = a.matmul(&b)?;
    println!("{}", c);
    Ok(())
}

• Device::Cpu é o padrão portátil; CUDA/Metal precisam de flags de recursos e drivers.
• candle espelha operações de tensor semelhantes ao PyTorch para modelos pequenos e pesquisa.
• Erros de forma aparecem como Result - valide as dimensões antecipadamente.

Relacionado: candle - módulos nn e treinamento

4. Executar um único passo de inferência ONNX

Alimente a entrada f32 e leia os logits.

use ndarray::Array;
use ort::{inputs, session::Session, value::Tensor};
 
fn main() -> ort::Result<()> {
    let session = Session::builder()?.commit_from_file("mnist.onnx")?;
    let input = Array::from_shape_vec((1, 1, 28, 28), vec![0.0f32; 784])?;
    let outputs = session.run(inputs![Tensor::from_array(input)?])?;
    let logits = outputs[0].try_extract_tensor::<f32>()?;
    println!("logits len: {}", logits.len());
    Ok(())
}

• A forma deve corresponder ao grafo ONNX - use convenções nhwd do script de exportação.
• Normalize os pixels da mesma forma que no treinamento, ou a precisão cairá.
• A dimensão do batch 1 ainda é importante para a compatibilidade do grafo.

5. Servir uma verificação de saúde com Axum

Encapsule a inferência via HTTP com um roteador mínimo.

use axum::{routing::get, Json, Router};
use serde_json::json;
 
#[tokio::main]
async fn main() {
    let app = Router::new().route("/health", get(|| async { Json(json!({"ok": true})) }));
    let listener = tokio::net::TcpListener::bind("127.0.0.1:8080").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

• Endpoints de saúde permitem que orquestradores sondem antes de enviar tráfego.
• Mantenha os pesos do modelo carregados uma vez na inicialização, não por requisição.
• Adicione /predict com corpos JSON validados por serde em seguida.

Relacionado: Servindo Modelos com Axum - manipuladores de produção

6. Serializar requisições de inferência com serde

Contratos de API tipados evitam erros de forma na borda.

use serde::{Deserialize, Serialize};
 
#[derive(Deserialize)]
struct PredictRequest {
    text: String,
}
 
#[derive(Serialize)]
struct PredictResponse {
    label: String,
    score: f32,
}

• Valide o comprimento máximo da string antes da tokenização para limitar a memória.
• Retorne scores com rótulos para observabilidade nos clientes.
• Espelhe a ordem das classes de treinamento nas strings de rótulo ou IDs.

7. Escolher burn para treinamento agnóstico de backend

Inicialize um backend explicitamente ao experimentar com burn.

use burn::backend::NdArray;
use burn::tensor::Tensor;
 
type B = NdArray<f32>;
 
fn main() {
    let device = <B as burn::backend::Backend>::Device::default();
    let t: Tensor<B, 2> = Tensor::from_data([[1.0, 2.0], [3.0, 4.0]], &device);
    let _ = t.matmul(t.transpose());
}

• burn troca os backends NdArray, WGPU e CUDA por trás de uma única API.
• Loops de treinamento vivem em Rust - nenhum notebook Python é necessário.
• Exporte para ONNX quando precisar de implantação multilíngue.

Relacionado: burn - loops de treinamento e checkpoints

Exemplos Intermediários

8. Batch de embeddings para recuperação

Calcule vetores para múltiplas strings em uma única passagem forward.

// Layout conceitual - combine tokenizers + ort ou modelo candle:
// 1. tokenizar com padding/truncamento
// 2. construir tensores de entrada [batch, seq]
// 3. executar sessão, pegar saída agrupada [batch, dim]

• O preenchimento alinha os comprimentos das sequências dentro de um batch na GPU.
• Normalize L2 os embeddings antes da busca por similaridade de cosseno.
• Armazene vetores em um índice do lado do Rust (veja a página de embeddings).

Relacionado: Embeddings e Busca Vetorial

9. Padrão de inferência de LLM local

Carregue pesos uma vez, transmita tokens com um crate de LLM Rust (bindings de llama.cpp ou mistral.rs).

// Estrutura de pseudocódigo para serviços de LLM local:
// let model = LlamaModel::load("weights.gguf", &params)?;
// let mut ctx = model.new_context()?;
// for token in ctx.generate(prompt, max_tokens) { print!("{token}") }

• Pesos GGUF são comuns para inferência local em CPU/GPU.
• O tamanho do cache KV determina a memória - limite max_seq_len por implantação.
• Execute a geração em um pool de threads dedicado para evitar bloquear os workers do Axum.

Relacionado: Inferência e Servir LLM

10. Flags de recursos de GPU

Habilite CUDA ou Metal apenas em builds de release que visam hosts de GPU.

[dependencies]
candle-core = { version = "0.8", features = ["cuda"] }

• Agentes CI apenas com CPU não devem exigir recursos de GPU.
• Use features Cargo separadas gpu vs builds padrão de CPU.
• Volte para CPU quando Device::cuda_if_available() falhar.

Relacionado: GPU e Aceleração

Versões da Stack: Esta página foi escrita para Rust 1.97.0 (edição 2024), Tokio 1.x, Axum 0.8, serde 1.0, sqlx 0.8, clap 4 e Polars 0.46+.