端侧 AI 的全景图：WASM、WebGPU 和 ONNX 的三角关系的深度剖析

端侧 AI 的全景图：WASM、WebGPU 和 ONNX 的三角关系的深度剖析

一、为什么 AI 必须走出云端

去年我在做一个离线翻译工具的需求调研时发现了一个有意思的数据：一款流行浏览器翻译插件的日活有 300 万，其中 60% 的网络请求被云端的翻译 API 消耗。这 60% 里有至少 40% 是完全可以本地处理的短文本翻译。

三个趋势在推动 AI 从云端走向端侧：

而端侧 AI 的落地依赖于三个技术的协作：WASM（容器）、WebGPU（加速器）、ONNX（模型格式）。这篇文章会厘清这个三角关系，并给出实际可用的技术选型框架。

二、ONNX：跨框架的模型交换格式

ONNX（Open Neural Network Exchange）解决的是模型格式碎片化的问题。PyTorch、TensorFlow、JAX 各有各的模型格式，ONNX 提供了一个中间表示层，让一个模型可以被多个推理引擎消费。

# Python 端：导出 ONNX 模型
import torch
import torch.onnx

# 定义一个简单的分类模型
class TextClassifier(torch.nn.Module):
def __init__(self, vocab_size: int, embed_dim: int, num_classes: int):
super().__init__()
self.embedding = torch.nn.Embedding(vocab_size, embed_dim)
self.fc = torch.nn.Linear(embed_dim, num_classes)

def forward(self, input_ids):
# 简单的嵌入 + 线性分类（示例用，实际用 BERT 更复杂）
x = self.embedding(input_ids).mean(dim=1) # 均值池化
return torch.softmax(self.fc(x), dim=-1)

model = TextClassifier(vocab_size=30000, embed_dim=128, num_classes=10)
model.eval()

# 导出为 ONNX 格式
dummy_input = torch.randint(0, 30000, (1, 32)) # batch=1, seq_len=32
torch.onnx.export(
model,
dummy_input,
"text_classifier.onnx",
input_names=["input_ids"],
output_names=["probs"],
dynamic_axes={
"input_ids": {0: "batch_size", 1: "seq_len"},
"probs": {0: "batch_size"},
},
opset_version=14,
)
print("模型已导出为 text_classifier.onnx")

ONNX 选型时注意两个关键点：

• opset version 兼容性——不是所有 ONNX Runtime 都支持最新的 opset。端侧一般用 opset 12-14，浏览器端 onnxruntime-web 对 opset 的兼容最好；
• 量化——FP32 模型在浏览器里跑不动的。导出时必须做 INT8 量化，内存占用能减到 1/4，推理速度提升 2-3 倍。

三、WASM：浏览器的通用运行时

WASM 解决了"浏览器能不能跑原生代码"的问题。Rust 通过 wasm-pack 编译到 WASM，再通过 JavaScript 加载调用。但对于 AI 推理，WASM 有天然的短板——纯 CPU 计算、SIMD 支持有限、不能直接访问 GPU。

// 用 wasm-bindgen 暴露 Rust 函数给 JavaScript 调用
use wasm_bindgen::prelude::*;
use wasm_bindgen_futures::JsFuture;
use web_sys::{Request, RequestInit, RequestMode, Response};

/// 在 Web Worker 中运行推理任务
/// 避免阻塞主线程导致页面卡顿
#[wasm_bindgen]
pub async fn run_inference_in_worker(
model_url: &str,
input_data: &[f32],
) -> Result<JsValue, JsValue> {
// 创建 Web Worker 实例
let worker = web_sys::Worker::new("./inference_worker.js")
.map_err(|e| JsValue::from_str(&format!("创建 Worker 失败: {:?}", e)))?;

// 在 Worker 中执行推理——不阻塞主线程
let promise = js_sys::Promise::new(&mut |resolve, _reject| {
// 向 Worker 发送数据和回调
let _ = worker.post_message(&JsValue::from_str("start_inference"));
resolve.call1(&JsValue::null(), &JsValue::from_str("inference_done")).unwrap();
});

JsFuture::from(promise).await
}

WASM 的定位是"浏览器里跑逻辑的最低门槛方案"。适合轻量级模型（< 50MB）、低频推理（用户手动触发而非实时）、不要求极致性能的场景。

四、WebGPU：端侧 GPU 推理的答案

WebGPU 是三角关系里最重要的角色——它让浏览器以近原生性能访问 GPU。和 WebGL 不同，WebGPU 是为计算而非渲染设计的，天然适合矩阵运算。

// 使用 burn-rs 的 WASM + WebGPU 后端
// burn-rs 抽象了 WASM 和原生 GPU 的差异
use burn::tensor::Tensor;
use burn::backend::wgpu::{WgpuDevice, Wgpu};
use burn::nn::{Linear, LinearConfig};
use burn::module::Module;
use burn::prelude::Backend;

/// 模型定义——burn-rs 的代码在后端切换时无需修改
#[derive(Module, Debug)]
pub struct SimpleTextModel<B: Backend> {
/// 线性层（在 CPU / WGPU / CUDA 上都能跑）
linear: Linear<B>,
}

impl<B: Backend> SimpleTextModel<B> {
pub fn new(device: &B::Device) -> Self {
let linear = LinearConfig::new(128, 10).init(device);
Self { linear }
}

pub fn forward(&self, input: Tensor<B, 2>) -> Tensor<B, 2> {
// 前向传播——后端的 GPU 加速自动生效
self.linear.forward(input)
}
}

// ===== 后端选择 =====

// CPU 后端（开发测试用）
// type MyBackend = burn::backend::NdArray;

// WASM 后端（浏览器 CPU）
// type MyBackend = burn::backend::Wgpu<f32, i32>;

// WebGPU 后端（浏览器 GPU——生产推荐）
type MyBackend = Wgpu<f32, i32>;

// CUDA 后端（原生桌面）
// type MyBackend = burn::backend::Cuda<f32, i32>;

WebGPU 选型的现实判断：

我们的 Whisper 小模型在 WASM（纯 CPU）上跑一次推理要 2.1 秒，切到 WebGPU 后降到 0.3 秒——快了 7 倍。但代价是 WebGPU 的初始化过程完全黑盒：设备创建失败时 Chrome 只返回 "validation error"，没有任何细节。最后靠 Google 搜到一个 GitHub issue 才知道是显存不够，给了 limits.maxStorageBufferBindingSize 一个更小的值才通过。

WebGPU 的错误信息比 WASM 差了一个数量级，排查时需要大量试错，建议先在小模型上跑通流程。

五、总结

WASM、WebGPU 和 ONNX 的三角关系可以这样理解：

• ONNX 决定"跑什么模型"——它是模型格式标准，解决了跨框架兼容问题；
• WASM 决定"在哪里跑"——它是浏览器环境的运行时容器，给了 Rust 代码进入浏览器的大门；
• WebGPU 决定"跑多快"——它是端侧 GPU 加速的入口，把推理速度从秒级降到十毫秒级。

我给开发者的技术选型建议：起步用 ONNX Runtime Web (WASM 后端)，模型跑通了再评估是否值得切 WebGPU。不要一上来就踩 WebGPU 的坑——GPU buffer 管理、shader 编译错误、不同浏览器的兼容差异，这些问题都会让你的开发效率腰斩。而 ONNX Runtime 的 WASM 后端只需要 npm install onnxruntime-web 就能开始用，成本极低。

端侧 AI 的真正爆发还需要一年左右——等主流浏览器对 WebGPU 的计算特性支持稳定，等 ONNX 生态的工具链再成熟一些。但现在开始投入学习和实验，正好能赶上第一波红利。