WebAssembly AI 插件：浏览器端 ONNX Runtime 推理与 Rust 模型封装

一、浏览器端推理的困境：为什么不能总是调用云端 API

Web 应用中越来越多的 AI 功能依赖云端 API：图像分类、文本摘要、语音识别。但每次调用都有 200-500ms 的网络延迟，加上 API 调用费用和隐私风险（用户数据上传到服务器）。浏览器端推理可以解决这些问题：零网络延迟、零 API 费用、数据不离开用户设备。WebAssembly 让 Rust 编写的推理引擎可以在浏览器中以接近原生的速度运行，ONNX Runtime Web 提供了 WASM 后端的推理能力。

graph TB
    A[AI 功能需求] --> B{推理方案}
    B --> C[云端 API]
    B --> D[浏览器端 WASM]

    C --> E[延迟: 200-500ms]
    C --> F[费用: 按调用计费]
    C --> G[隐私: 数据上传]

    D --> H[延迟: 50-200ms<br/>取决于模型大小]
    D --> I[费用: 零]
    D --> J[隐私: 本地推理]

    D --> K[Rust 训练/导出模型]
    K --> L[ONNX 格式导出]
    L --> M[wasm-pack 编译]
    M --> N[浏览器加载 .wasm]
    N --> O[ONNX Runtime Web 推理]

二、Rust → ONNX → WASM 推理管线的底层机制

2.1 ONNX Runtime Web 的执行提供者

ONNX Runtime Web 支持三种执行后端：WASM（CPU 通用）、WebGL（GPU 加速）、WebGPU（下一代 GPU API）。WASM 后端兼容性最好但速度最慢，WebGPU 后端速度最快但浏览器支持有限。

graph LR
    A[ONNX 模型文件] --> B{执行提供者}
    B --> C[WASM CPU<br/>兼容性: 全平台<br/>速度: 基线]
    B --> D[WebGL GPU<br/>兼容性: 主流浏览器<br/>速度: 2-5x]
    B --> E[WebGPU<br/>兼容性: Chrome 113+<br/>速度: 5-20x]

    C --> F[适合: 小模型<br/>文本分类/NER]
    D --> G[适合: 中模型<br/>图像分类]
    E --> H[适合: 大模型<br/>目标检测/分割]

2.2 Rust 模型封装与 wasm-bindgen 桥接

Rust 代码通过 wasm-bindgen 暴露为 JavaScript API，负责模型加载、预处理、推理调度和后处理。ONNX Runtime Web 在 JavaScript 侧运行，Rust 侧通过 JS 互操作调用推理接口。

2.3 模型量化与体积优化

浏览器端模型体积直接影响加载时间。INT8 量化可将模型体积压缩 4 倍（FP32 → INT8），精度损失通常小于 1%。对于移动端场景，还可使用知识蒸馏训练更小的学生模型。

三、生产级代码实现与最佳实践

3.1 Rust 侧推理引擎封装

use wasm_bindgen::prelude::*;
use serde::{Deserialize, Serialize};

/// 推理结果
#[derive(Serialize, Deserialize)]
pub struct InferenceResult {
    pub label: String,
    pub confidence: f32,
    pub latency_ms: f64,
}

/// 图像分类器（Rust 侧封装，JS 互操作）
#[wasm_bindgen]
pub struct ImageClassifier {
    model_bytes: Vec<u8>,
    labels: Vec<String>,
    input_size: (usize, usize),
}

#[wasm_bindgen]
impl ImageClassifier {
    /// 从 ArrayBuffer 加载 ONNX 模型
    #[wasm_bindgen(constructor)]
    pub fn new(model_bytes: &[u8], labels_json: &str) -> Result<ImageClassifier, JsValue> {
        let labels: Vec<String> = serde_json::from_str(labels_json)
            .map_err(|e| JsValue::from_str(&format!("标签解析失败: {}", e)))?;

        Ok(ImageClassifier {
            model_bytes: model_bytes.to_vec(),
            labels,
            input_size: (224, 224), // MobileNet 默认输入尺寸
        })
    }

    /// 对图像数据进行分类
    /// pixels: RGBA 格式的 Uint8Array
    pub async fn classify(&self, pixels: &[u8], width: usize, height: usize) -> Result<JsValue, JsValue> {
        let start = js_sys::Date::now();

        // 1. 图像预处理：缩放 + 归一化
        let input_tensor = self.preprocess(pixels, width, height);

        // 2. 调用 ONNX Runtime Web 推理（通过 JS 互操作）
        let output = self.run_inference(&input_tensor).await?;

        // 3. 后处理：Softmax + Top-1
        let result = self.postprocess(&output);

        let latency = js_sys::Date::now() - start;
        let result_with_latency = InferenceResult {
            label: result.0,
            confidence: result.1,
            latency_ms: latency,
        };

        // 序列化为 JS 对象
        serde_wasm_bindgen::to_value(&result_with_latency)
            .map_err(|e| JsValue::from_str(&format!("序列化失败: {}", e)))
    }

    /// 图像预处理：缩放到 224x224，归一化到 [0, 1]，转为 NCHW 格式
    fn preprocess(&self, pixels: &[u8], width: usize, height: usize) -> Vec<f32> {
        let (target_w, target_h) = self.input_size;
        let channels = 3; // RGB
        let mut tensor = vec![0.0f32; 1 * channels * target_w * target_h];

        // ImageNet 归一化参数
        let mean = [0.485, 0.456, 0.406];
        let std = [0.229, 0.224, 0.225];

        for y in 0..target_h {
            for x in 0..target_w {
                // 双线性插值采样
                let src_x = (x as f32 * width as f32 / target_w as f32) as usize;
                let src_y = (y as f32 * height as f32 / target_h as f32) as usize;
                let src_idx = (src_y * width + src_x) * 4; // RGBA

                if src_idx + 2 < pixels.len() {
                    // NCHW 布局: [batch, channel, height, width]
                    for c in 0..channels {
                        let pixel_value = pixels[src_idx + c] as f32 / 255.0;
                        let normalized = (pixel_value - mean[c]) / std[c];
                        let dst_idx = c * target_w * target_h + y * target_w + x;
                        tensor[dst_idx] = normalized;
                    }
                }
            }
        }

        tensor
    }

    /// 调用 ONNX Runtime Web 推理
    async fn run_inference(&self, input_tensor: &[f32]) -> Result<Vec<f32>, JsValue> {
        // 通过 js_sys 调用 ONNX Runtime Web 的 JavaScript API
        let js_result = js_sys::eval("window.ortSessionRun").unwrap();
        let run_fn: js_sys::Function = js_result.into();

        // 将输入张量转为 JS Float32Array
        let input_array = js_sys::Float32Array::new_with_length(input_tensor.len() as u32);
        input_array.copy_from(input_tensor);

        let promise = run_fn.call1(&JsValue::NULL, &input_array)?;
        let result = wasm_bindgen_futures::JsFuture::from(
            js_sys::Promise::resolve(&promise)
        ).await?;

        // 解析输出张量
        let output_array: js_sys::Float32Array = result.into();
        let mut output = vec![0.0f32; output_array.length() as usize];
        output_array.copy_to(&mut output);
        Ok(output)
    }

    /// Softmax + Top-1 后处理
    fn postprocess(&self, logits: &[f32]) -> (String, f32) {
        // 数值稳定的 Softmax
        let max_logit = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
        let exp_sum: f32 = logits.iter()
            .map(|&x| (x - max_logit).exp())
            .sum();
        let probs: Vec<f32> = logits.iter()
            .map(|&x| (x - max_logit).exp() / exp_sum)
            .collect();

        // Top-1
        let (best_idx, &best_prob) = probs.iter()
            .enumerate()
            .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
            .unwrap();

        let label = self.labels.get(best_idx)
            .cloned()
            .unwrap_or_else(|| format!("unknown_{}", best_idx));

        (label, best_prob)
    }
}

3.2 JavaScript 侧 ONNX Runtime 初始化

// ort-init.js — ONNX Runtime Web 初始化与推理函数
import { InferenceSession, Tensor } from 'onnxruntime-web';

let session = null;

// 初始化推理会话（页面加载时调用）
export async function initSession(modelUrl) {
    const opt = {
        executionProviders: ['webgpu', 'wasm'],  // 优先 WebGPU，回退 WASM
        graphOptimizationLevel: 'all'
    };
    session = await InferenceSession.create(modelUrl, opt);
    console.log(`ONNX Runtime 会话已初始化，后端: ${session.handler.backend}`);
}

// 供 Rust WASM 调用的推理函数
window.ortSessionRun = async function(inputFloat32Array) {
    if (!session) throw new Error('推理会话未初始化');

    // 构造输入张量: [1, 3, 224, 224]
    const inputTensor = new Tensor('float32', inputFloat32Array, [1, 3, 224, 224]);
    const feeds = { [session.inputNames[0]]: inputTensor };

    const results = await session.run(feeds);
    const outputTensor = results[session.outputNames[0]];

    // 返回 Float32Array 给 Rust
    return outputTensor.data;
};

3.3 模型量化与构建管线

# Cargo.toml — WASM 构建配置
[package]
name = "wasm-ai-plugin"
version = "0.1.0"
edition = "2021"

[lib]
crate-type = ["cdylib"]

[dependencies]
wasm-bindgen = "0.2"
wasm-bindgen-futures = "0.4"
js-sys = "0.3"
serde = { version = "1", features = ["derive"] }
serde_json = "1"
serde-wasm-bindgen = "0.6"

[profile.release]
opt-level = "z"      # 优化体积
lto = true           # 链接时优化
strip = true         # 去除调试信息

# 构建命令
# 1. Python 侧模型量化
python -m onnxruntime.quantization.quantize_static \
    --model_input mobilenet_v2.onnx \
    --model_output mobilenet_v2_int8.onnx \
    --quant_format QDQ \
    --per_channel \
    --weight_type int8

# 2. Rust 编译为 WASM
wasm-pack build --target web --release

# 3. 输出文件
# pkg/wasm_ai_plugin.js       — JS 胶水代码
# pkg/wasm_ai_plugin_bg.wasm  — WASM 二进制
# pkg/wasm_ai_plugin.d.ts     — TypeScript 类型

四、浏览器端推理的架构权衡

4.1 推理速度 vs 模型精度

量化方案	模型体积	推理速度 (WASM)	Top-1 精度损失
FP32 原始	14MB	基线	0%
FP16 量化	7MB	~1.2x	< 0.1%
INT8 静态量化	3.5MB	~1.5x	0.5-2%
INT8 + 蒸馏小模型	1.8MB	~3x	2-5%

4.2 首次加载 vs 缓存复用

WASM 文件和 ONNX 模型首次加载需要下载，MobileNet V2 INT8 总计约 5MB。使用 HTTP Cache-Control 和 Service Worker 缓存后，二次访问加载时间可降至 50ms 以内。

4.3 适用边界与禁用场景

适用场景：

• 图像分类、文本分类等轻量推理任务
• 隐私敏感场景（医疗影像、个人照片分析）
• 离线可用的 Web 应用

禁用场景：

• 大语言模型推理（模型体积 > 1GB，浏览器内存不足）
• 实时视频流处理（WASM 后端帧率不足）
• 需要高精度数值计算的任务（INT8 量化精度不够）

五、总结

浏览器端 AI 推理的核心价值是零延迟和零隐私泄露。Rust 通过 wasm-bindgen 将预处理和后处理逻辑编译为 WASM，与 ONNX Runtime Web 的推理能力组合，形成完整的端侧推理管线。INT8 量化是浏览器场景的必选项——3.5MB 的模型比 14MB 的模型加载快 4 倍，精度损失在可接受范围内。WebGPU 后端是未来的性能突破口，但当前兼容性限制了生产使用。对于轻量推理任务，WASM + ONNX Runtime Web 已经是可用的生产方案。