在浏览器里跑AI模型：WebAssembly插件开发实战与推理优化

一、当AI推理走出服务器：浏览器端的机遇与挑战

去年我第一次在浏览器里跑通一个文本分类模型时，那种兴奋感至今记得。不需要服务器，不需要GPU，打开网页就能推理。这意味着什么？意味着AI能力可以真正分发到每个用户手中，零部署成本，零隐私泄露风险。

但兴奋过后，现实的问题接踵而至。模型太大加载慢，推理延迟高，内存占用吓人。一个50MB的模型在浏览器里要加载好几秒，推理一次要等半秒。这在生产环境中是不可接受的。

WebAssembly给这些问题提供了一个解法。WASM的接近原生执行速度、沙箱安全性和跨语言编译能力，让它成为浏览器端AI推理的理想载体。这篇文章记录我用Rust+WASM开发AI推理插件的完整过程，包括踩过的坑和找到的优化路径。

二、WASM AI插件的技术栈与运行机制

先理清整体技术链路。一个WASM AI插件从源码到运行，要经历这些阶段：

graph LR
    A[Rust源码] --> B[cargo build --target wasm32-unknown-unknown]
    B --> W[.wasm文件]
    W --> C[wasm-bindgen处理]
    C --> D[生成JS胶水代码]
    D --> E[打包工具集成]
    E --> F[浏览器加载]
    F --> G[WASM运行时执行]
    H[ONNX模型文件] --> I[编译为WASM兼容格式]
    I --> F
    G --> J[推理结果返回JS]
    J --> K[UI渲染]

核心组件有三个层次：

模型层： ONNX格式是当前最通用的选择。PyTorch和TensorFlow都能导出ONNX，而tract库可以在Rust中直接加载和执行ONNX模型，不需要依赖外部运行时。

推理层： 这是WASM插件的核心。用Rust实现推理逻辑，编译为WASM模块。关键在于避免使用Rust标准库中WASM不支持的特性（如文件IO、线程、网络）。

交互层： 通过wasm-bindgen暴露Rust函数给JavaScript调用。模型输入输出都通过这个桥梁传递。

一个容易忽略的点：WASM线性内存的限制。当前WASM规范中，线性内存最大4GB。对于大多数轻量模型够用，但如果你想做图像生成，这个限制就是硬约束。

三、完整实现：一个文本分类WASM插件

下面是一个完整的文本情感分类插件，从Rust代码到浏览器调用。

首先是Cargo.toml配置：

[package]
name = "wasm-sentiment"
version = "0.1.0"
edition = "2021"

[lib]
crate-type = ["cdrlib"]  # 编译为WASM动态库

[dependencies]
wasm-bindgen = "0.2"
tract-onnx = "0.21"
serde = { version = "1", features = ["derive"] }
serde_json = "1"

# 为什么用console_log？因为WASM里println!不工作
# 必须通过JS的console输出日志
console_log = "1"
log = "0.4"

[profile.release]
# 优化WASM体积的关键配置
opt-level = "z"     # 体积优先优化
lto = true          # 链接时优化，去掉未使用代码
strip = true        # 移除调试信息

核心推理代码：

use wasm_bindgen::prelude::*;
use tract_onnx::prelude::*;

/// 情感分析模型包装器
/// 为什么用struct而不是直接暴露函数？
/// 因为模型加载是重量级操作，应该只做一次
#[wasm_bindgen]
pub struct SentimentModel {
    model: RunnableModel<TypedFact, Box<dyn TypedOp>, Graph<TypedFact, Box<dyn TypedOp>>>,
    tokenizer: SimpleTokenizer,
}

/// 简易分词器：将文本转为模型输入的token ID序列
/// 生产环境建议用tokenizers库，这里为了减少WASM体积用简化版
struct SimpleTokenizer {
    vocab: std::collections::HashMap<String, usize>,
}

impl SimpleTokenizer {
    fn new() -> Self {
        // 硬编码一个小型词表，实际项目从模型元数据加载
        let vocab: std::collections::HashMap<String, usize> = [
            ("好", 1), ("棒", 2), ("差", 3), ("烂", 4),
            ("不错", 5), ("一般", 6), ("喜欢", 7), ("讨厌", 8),
            ("[PAD]", 0), ("[CLS]", 101), ("[SEP]", 102),
        ].iter().map(|(k, v)| (k.to_string(), *v)).collect();

        Self { vocab }
    }

    fn tokenize(&self, text: &str) -> Vec<usize> {
        let mut ids = vec![101]; // [CLS]
        // 简单的逐字符+双字符匹配
        let chars: Vec<char> = text.chars().collect();
        let mut i = 0;
        while i < chars.len() {
            if i + 1 < chars.len() {
                let bigram = format!("{}{}", chars[i], chars[i + 1]);
                if let Some(&id) = self.vocab.get(&bigram) {
                    ids.push(id);
                    i += 2;
                    continue;
                }
            }
            let ch = chars[i].to_string();
            ids.push(*self.vocab.get(&ch).unwrap_or(&0));
            i += 1;
        }
        ids.push(102); // [SEP]
        ids
    }
}

#[wasm_bindgen]
impl SentimentModel {
    /// 从字节数组加载模型
    /// 为什么用字节数组而不是文件路径？
    /// 因为WASM环境没有文件系统，模型数据通过JS传入
    #[wasm_bindgen(constructor)]
    pub fn new(model_bytes: &[u8]) -> Result<SentimentModel, JsValue> {
        // tract支持从内存加载ONNX模型
        let model = tract_onnx::onnx()
            .model_for_read(&mut &model_bytes[..])
            .map_err(|e| JsValue::from_str(&format!("模型加载失败: {}", e)))?
            .into_optimized()
            .map_err(|e| JsValue::from_str(&format!("模型优化失败: {}", e)))?
            .into_runnable()
            .map_err(|e| JsValue::from_str(&format!("模型编译失败: {}", e)))?;

        Ok(SentimentModel {
            model,
            tokenizer: SimpleTokenizer::new(),
        })
    }

    /// 执行推理，返回情感分数
    pub fn predict(&self, text: &str) -> Result<f32, JsValue> {
        let token_ids = self.tokenizer.tokenize(text);

        // 构建输入张量
        // 为什么用i64？因为ONNX模型通常期望int64输入
        let input: Tensor = tract_ndarray::Array2::from_shape_vec(
            (1, token_ids.len()),
            token_ids.into_iter().map(|id| id as i64).collect::<Vec<_>>(),
        ).unwrap().into();

        // 执行推理
        let result = self.model.run(tvec!(input.into()))
            .map_err(|e| JsValue::from_str(&format!("推理失败: {}", e)))?;

        // 解析输出：取[CLS]位置的logits
        let logits = result[0].to_array_view::<f32>()
            .map_err(|e| JsValue::from_str(&format!("输出解析失败: {}", e)))?;

        // softmax得到概率
        let positive_score = logits[[0, 1]].exp();
        let negative_score = logits[[0, 0]].exp();
        let total = positive_score + negative_score;

        Ok(positive_score / total)
    }
}

JavaScript端调用：

import init, { SentimentModel } from './pkg/wasm_sentiment.js';

async function runInference() {
    // 初始化WASM模块
    await init();

    // 加载模型（实际项目中模型文件通过fetch获取）
    const modelResponse = await fetch('/models/sentiment.onnx');
    const modelBytes = await modelResponse.arrayBuffer();

    // 创建模型实例（只做一次，后续复用）
    const model = new SentimentModel(new Uint8Array(modelBytes));

    // 推理
    const score = model.predict("这个产品真的很棒，推荐购买");
    console.log(`正面情感概率: ${(score * 100).toFixed(1)}%`);
}

四、性能优化的取舍与边界

开发过程中，我做了几个关键优化，每个都有代价。

模型量化。 将FP32权重转为INT8，模型体积缩小4倍，推理速度提升2-3倍。代价是精度损失，情感分类这类任务影响不大，但NLU任务可能明显掉点。我的做法是：先量化，再在验证集上评估，掉点超过2%就不用量化。

WASM模块体积。 初始编译出来8MB，经过opt-level="z"、LTO、strip后降到1.2MB。但tract库本身很大，如果只需要推理不需要加载ONNX，可以用tract-core替代tract-onnx，再省几百KB。另一个思路是用wasm-opt做后处理优化，能再压缩10-15%。

内存管理。 Rust侧分配的内存不会自动被JS垃圾回收。必须在Rust侧实现Drop，或者在JS侧手动调用释放函数。我踩过这个坑：长时间运行后浏览器内存持续增长。解决方案是给SentimentModel加free方法：

#[wasm_bindgen]
impl SentimentModel {
    /// 必须手动释放，否则WASM线性内存泄漏
    pub fn free(&mut self) {
        // Drop会自动执行，这里显式标记
    }
}

SIMD加速。 WASM SIMD支持128位向量运算，对矩阵乘法加速明显。但需要浏览器支持，Safari直到2023年才默认启用。如果你的用户群体包含老浏览器，需要做降级方案。

五、总结

WebAssembly让浏览器端AI推理从"概念验证"走向"生产可用"。Rust+WASM的组合提供了性能和安全性的平衡点。但这条路并不平坦：模型体积、内存管理、浏览器兼容性都是需要持续优化的方向。

我的建议是：从轻量模型开始，先跑通链路，再逐步优化。不要一上来就想在浏览器里跑LLM，先从分类、检测这类小模型入手，积累经验后再挑战更大的场景。技术探索需要耐心，也需要务实。