WASM + AI：浏览器端推理的架构设计与落地实践

WASM + AI：浏览器端推理的架构设计与落地实践

一、AI 推理进浏览器：不是炫技，是刚需

把 AI 模型跑在浏览器里，听起来像技术演示，但在实际业务中有明确的驱动力。数据隐私是第一位的：医疗影像分析、金融文档处理，这些场景下数据不能离开用户设备。离线可用是第二位的：弱网环境、飞行模式，云端 API 不可用时本地推理是唯一选择。低延迟是第三位的：实时图像滤镜、语音识别，往返服务器的延迟不可接受。

WebAssembly 让这件事变得可行。它提供了接近原生的执行速度，沙箱化的安全模型，以及跨浏览器的一致运行时。但把一个训练好的模型变成浏览器里能跑的 WASM 模块，中间要解决的问题远不止"编译一下"这么简单。

模型体积是第一个拦路虎。一个 ResNet-50 模型的 ONNX 文件约 100MB，浏览器加载这个体积的 WASM 模块几乎不可接受。量化、剪枝、知识蒸馏——模型压缩是绕不开的前置步骤。推理性能是第二个问题。WASM 目前不支持 SIMD 在所有浏览器上的完整实现（Safari 的支持滞后），这直接影响矩阵运算的吞吐。内存管理是第三个问题。WASM 线性内存是固定大小的，模型权重和中间张量共享这块内存，规划不当就会 OOM。

二、WASM AI 推理的端到端架构

一个完整的浏览器端 AI 推理系统，涉及从模型训练到浏览器执行的完整链路。

graph LR
    A[训练好的模型 PyTorch/TF] --> B[模型导出 ONNX]
    B --> C[模型优化 量化/剪枝]
    C --> D[编译为 WASM Emscripten/wasm-pack]
    D --> E[Web 运行时加载]
    E --> F[前端预处理]
    F --> G[WASM 推理执行]
    G --> H[后处理与渲染]

    subgraph 浏览器端
        E
        F
        G
        H
    end

    subgraph 构建时
        A
        B
        C
        D
    end

构建时和运行时的分离是关键。构建时负责模型压缩和 WASM 编译，运行时只做加载和推理。这种分离意味着你可以在 CI/CD 中完成所有重计算，浏览器里只执行轻量的推理逻辑。

WASM 推理引擎的选择目前主要有两个方向：一是将现有的 C/C++ 推理框架（如 ONNX Runtime、TensorFlow Lite）编译为 WASM，二是用 Rust 编写推理逻辑并通过 wasm-pack 编译。前者兼容性好但产物体积大，后者灵活但需要自己实现算子。

三、用 Rust + wasm-pack 构建浏览器端图像分类器

以下代码展示了一个完整的 Rust → WASM 图像分类推理模块：

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

/// 分类结果
#[derive(Serialize, Deserialize)]
pub struct ClassResult {
    /// 类别索引
    pub class_id: usize,
    /// 置信度
    pub confidence: f32,
    /// 类别标签
    pub label: String,
}

/// 图像分类推理器
#[wasm_bindgen]
pub struct ImageClassifier {
    /// 模型权重（量化后的 u8 数组）
    weights: Vec<u8>,
    /// 输入尺寸
    input_size: usize,
    /// 类别标签列表
    labels: Vec<String>,
}

#[wasm_bindgen]
impl ImageClassifier {
    /// 从 WASM 内存中加载模型权重
    #[wasm_bindgen(constructor)]
    pub fn new(weights: &[u8], input_size: usize, labels: Vec<JsValue>) -> Result<ImageClassifier, JsValue> {
        let label_strings: Vec<String> = labels
            .iter()
            .filter_map(|v| v.as_string())
            .collect();

        if label_strings.is_empty() {
            return Err(JsValue::from_str("标签列表不能为空"));
        }

        Ok(ImageClassifier {
            weights: weights.to_vec(),
            input_size,
            labels: label_strings,
        })
    }

    /// 执行推理，接收预处理后的像素数据
    pub fn predict(&self, pixels: &[f32]) -> Result<JsValue, JsValue> {
        let expected_len = self.input_size * self.input_size * 3;
        if pixels.len() != expected_len {
            return Err(JsValue::from_str(&format!(
                "输入长度不匹配：期望 {}，实际 {}",
                expected_len,
                pixels.len()
            )));
        }

        // 执行简化的推理逻辑（实际应使用量化权重做矩阵运算）
        let scores = self.forward(pixels);

        // 取 Top-3 结果
        let mut indexed: Vec<(usize, f32)> = scores
            .iter()
            .enumerate()
            .map(|(i, &s)| (i, s))
            .collect();
        indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
        indexed.truncate(3);

        let results: Vec<ClassResult> = indexed
            .into_iter()
            .map(|(id, conf)| ClassResult {
                class_id: id,
                confidence: conf,
                label: self.labels.get(id)
                    .cloned()
                    .unwrap_or_else(|| format!("unknown_{}", id)),
            })
            .collect();

        // 序列化为 JSON 返回给 JS
        serde_wasm_bindgen::to_value(&results)
            .map_err(|e| JsValue::from_str(&e.to_string()))
    }

    /// 前向传播（简化实现，生产环境应替换为真正的量化推理）
    fn forward(&self, pixels: &[f32]) -> Vec<f32> {
        // 这里应是实际的量化矩阵运算
        // 简化示例：用全局平均池化模拟
        let num_classes = self.labels.len();
        let chunk_size = pixels.len() / num_classes;
        (0..num_classes)
            .map(|i| {
                let start = i * chunk_size;
                let end = (start + chunk_size).min(pixels.len());
                let sum: f32 = pixels[start..end].iter().sum();
                sum / chunk_size.max(1) as f32
            })
            .collect()
    }
}

对应的 JavaScript 调用代码：

import init, { ImageClassifier } from './pkg/image_classifier.js';

async function runInference(imageElement) {
    await init();

    // 从 Canvas 获取像素数据并预处理
    const canvas = document.createElement('canvas');
    canvas.width = 224;
    canvas.height = 224;
    const ctx = canvas.getContext('2d');
    ctx.drawImage(imageElement, 0, 0, 224, 224);
    const imageData = ctx.getImageData(0, 0, 224, 224);

    // 归一化到 [0, 1]
    const pixels = new Float32Array(224 * 224 * 3);
    for (let i = 0; i < 224 * 224; i++) {
        pixels[i * 3] = imageData.data[i * 4] / 255.0;
        pixels[i * 3 + 1] = imageData.data[i * 4 + 1] / 255.0;
        pixels[i * 3 + 2] = imageData.data[i * 4 + 2] / 255.0;
    }

    // 加载模型权重
    const weightsResponse = await fetch('models/quantized_weights.bin');
    const weights = new Uint8Array(await weightsResponse.arrayBuffer());

    const labels = ['cat', 'dog', 'bird', 'fish', 'car'];
    const classifier = new ImageClassifier(weights, 224, labels);
    const results = classifier.predict(pixels);

    console.log('分类结果：', results);
}

四、WASM AI 推理的边界与架构妥协

模型体积的硬约束：WASM 模块的加载时间直接影响用户体验。一个经验值是：WASM 文件超过 5MB 时，首次加载时间在 3G 网络下会超过 3 秒。这意味着大模型必须量化到 Int8 甚至 Int4，同时接受精度损失。量化不是免费的，分类任务的 Top-1 精度通常下降 1-3%，检测任务可能下降更多。

SIMD 支持的碎片化：WASM SIMD 在 Chrome 和 Firefox 中已稳定支持，但 Safari 的支持进度滞后。如果你的目标用户包含 iOS Safari，就不能依赖 SIMD 加速，推理性能可能下降 2-4 倍。一个务实的做法是编译两个版本的 WASM：带 SIMD 的和不带 SIMD 的，运行时检测支持情况后加载对应版本。

线程模型的限制：WASM 多线程依赖 SharedArrayBuffer，而 SharedArrayBuffer 要求页面设置特定的 COOP/COEP 安全头。很多现有站点无法满足这个要求，导致 WASM 多线程不可用。单线程推理的性能天花板明显，尤其是大语言模型的推理。

内存管理的坑：WASM 线性内存默认是 256MB 封顶（可通过配置扩展），但浏览器对单个 WASM 实例的内存有不同限制。Chrome 相对宽松，Safari 更严格。模型权重、输入张量、中间激活值共享这块内存，需要仔细规划。一个常见的做法是将权重放在 JS 侧的 ArrayBuffer 中，推理时通过 WebAssembly.Memory 的视图传递，避免重复拷贝。

五、总结

WASM AI 推理在数据隐私、离线可用和低延迟场景下有明确价值。架构上，构建时负责模型压缩和 WASM 编译，运行时只做加载和推理。Rust + wasm-pack 是当前最灵活的技术路线，但需要自行实现推理算子。主要瓶颈在于模型体积、SIMD 支持碎片化、线程模型受限和内存管理。落地时建议先做模型量化到 Int8，控制 WASM 产物在 5MB 以内，编译带/不带 SIMD 的双版本，并在运行时检测特性支持。WASM AI 推理不是万能方案，但在特定场景下，它是浏览器端唯一可行的选择。