端侧推理：利用 WebGPU 执行 ONNX 模型的实践
随着大模型与深度学习在前端场景的纵深渗透，许多传统依赖服务端 API 的 AI 特性（如文本向量嵌入计算、实时图像分割、敏感内容本地过滤、手势与人脸追踪）正加速向浏览器端侧迁移。
传统的 WASM（WebAssembly）推理受限于单指令多数据流（SIMD）宽度与 CPU 多核利用率，在面对数千万参数的神经网络时延迟往往高达数秒。
随着 WebGPU 在现代主流浏览器中的全面铺开，配合 ONNX Runtime Web（onnxruntime-web/webgpu），前端开发者可以直接调用 GPU 的并行着色器内核来执行深度学习张量计算，将端侧推理性能提升数倍至数十倍。
架构拓扑：端侧推理流水线
在端侧部署 AI 推理，整体架构由模型加载、内存预分配、张量转换与着色器流水线组成：
[ ONNX 量化模型文件 (.onnx 15MB~50MB) ]
│
▼ (Cache Storage / IndexedDB 离线缓存)
┌─────────────────────────────────────────────────────────┐
│ ONNX Runtime Web (WebGPU EP) │
│ ├─ 1. Graph Optimizer (算子融合、常数折叠) │
│ ├─ 2. WGSL JIT (根据算子动态生成 WebGPU Compute Shader) │
│ └─ 3. WebGPU Storage Buffer (显存张量驻留池) │
└──────────────────────────┬──────────────────────────────┘
│
[ 端侧推理 Execute ]
│ (零网络请求，延迟 < 15ms)
▼
[ 实时特征向量 / 掩膜上屏 ]
环境初始化与模型装配实战
在项目中引入 onnxruntime-web 时，需明确指定后端执行提供商（Execution Provider）为 webgpu，并开启显存共享与预热机制：
import * as ort from 'onnxruntime-web/webgpu';
// 配置 WASM 与 WebGPU 依赖文件路径（生产环境应部署至本地 CDN）
ort.env.wasm.numThreads = navigator.hardwareConcurrency || 4;
ort.env.wasm.simd = true;
export interface InferenceEngineConfig {
modelUrl: string;
executionProvider?: 'webgpu' | 'wasm';
}
export class OnnxWebGPUEngine {
private session: ort.InferenceSession | null = null;
private isWarmedUp: boolean = false;
public async initialize(config: InferenceEngineConfig): Promise<void> {
const ep = config.executionProvider || 'webgpu';
const sessionOptions: ort.InferenceSession.SessionOptions = {
executionProviders: [
{
name: ep,
deviceType: 'gpu',
powerPreference: 'high-performance'
}
],
graphOptimizationLevel: 'all',
enableCpuMemArena: true,
enableMemPattern: true
};
console.log(`[ONNX-WebGPU] 正在加载模型: ${config.modelUrl}`);
this.session = await ort.InferenceSession.create(config.modelUrl, sessionOptions);
console.log('[ONNX-WebGPU] 模型加载与计算图编译完成');
// 必须进行预热推理（Warmup），触发 WebGPU Compute Pipeline 的 JIT 编译
await this.warmup();
}
/**
* 预热执行：首次推理包含着色器编译耗时，需在后台完成
*/
private async warmup() {
if (!this.session) return;
try {
// 构造全零虚拟输入张量（以典型 1x3x224x224 图像模型为例）
const dummyData = new Float32Array(1 * 3 * 224 * 224);
const dummyTensor = new ort.Tensor('float32', dummyData, [1, 3, 224, 224]);
const inputName = this.session.inputNames[0];
const feeds: Record<string, ort.Tensor> = { [inputName]: dummyTensor };
const start = performance.now();
await this.session.run(feeds);
console.log(`[ONNX-WebGPU] 预热编译耗时: ${(performance.now() – start).toFixed(2)}ms`);
this.isWarmedUp = true;
} catch (e) {
console.warn('[ONNX-WebGPU] 预热失败，将在首次正式请求时编译', e);
}
}
/**
* 执行推理
*/
public async runInference(inputData: Float32Array, dims: number[]): Promise<Float32Array> {
if (!this.session) {
throw new Error('Inference session is not initialized');
}
const inputTensor = new ort.Tensor('float32', inputData, dims);
const inputName = this.session.inputNames[0];
const outputName = this.session.outputNames[0];
const feeds: Record<string, ort.Tensor> = { [inputName]: inputTensor };
const start = performance.now();
const results = await this.session.run(feeds);
const cost = performance.now() – start;
console.log(`[ONNX-WebGPU] 推理耗时: ${cost.toFixed(2)}ms`);
const outputTensor = results[outputName];
return outputTensor.data as Float32Array;
}
public async release() {
if (this.session) {
await this.session.release();
this.session = null;
}
}
}
生产级量化与体积优化
直接将 100MB+ 的 FP32 原生 ONNX 模型丢给浏览器是不切实际的。在生产部署前必须完成三道工序：
故障降级底线
虽然主流浏览器（Chrome 113+、Edge、Safari 17+）已逐步开启 WebGPU 支持，但低端核显、移动端老旧设备或 Linux 驱动异常环境仍可能发生 navigator.gpu 为空或着色器编译崩溃。
因此，端侧推理引擎必须设计严格的降级链条：WebGPU (极速硬件加速) $\\to$ WASM + SIMD (多核 CPU 托底) $\\to$ 服务端 API 兜底 (网络回退)。
端侧算力的觉醒，让前端工程师不再只是被动消费服务端数据的“画皮者”，而是能够将万亿次浮点运算凝聚在用户掌中方寸之间的全栈计算掌控者。
