我花3个月写了个浏览器端 ML 框架，性能提升 3 倍
GitHub: github.com/s-zx/edgeFl… 求 ⭐ Star 支持
前言
最近 AI 很火，浏览器端跑模型的需求也越来越多。我在做一个项目时用了 transformers.js，但遇到了一些让我很难受的问题：
- 只能串行推理 - 我需要同时跑 4 个模型，结果发现只能一个一个来，180ms 的延迟用户体验很差
- 内存泄漏 - 跑了几个小时后页面就崩了，查了半天发现是 Tensor 没释放
- 包太大了 - 光 ONNX Runtime 就 2MB+，首屏加载慢得要命
- 下载体验差 - 500MB 的模型下到一半断了，只能从头来
于是我决定自己造个轮子 —— edgeFlow.js。
效果对比
先上个数据：
| 指标 | transformers.js | edgeFlow.js |
|---|---|---|
| 4模型并发 | 180ms (串行) | 52ms (并行) |
| 包体积 | ~2-5MB | <500KB |
| 内存管理 | 基础 | RAII + 泄漏检测 |
| 下载支持 | 普通 | 分片 + 断点续传 |
（是的，性能提升 3 倍不是标题党）
核心设计
1. 并发调度器
这是整个框架的核心。设计目标：
- 支持多模型并行执行
- 优先级队列（紧急任务插队）
- 模型级别隔离（避免竞态）
- 可选的批处理（合并小请求）
关键代码：
class InferenceScheduler {
private queues: Map<string, PriorityQueue<Task>> = new Map();
private runningTasks: Map<string, Set<string>> = new Map();
schedule(modelId, executor, priority) {
const task = new Task(modelId, executor, priority);
this.getQueue(modelId).enqueue(task);
this.processQueue();
return task;
}
private processQueue() {
for (const [modelId, queue] of this.queues) {
while (!queue.isEmpty() && this.canStartTask(modelId)) {
const task = queue.dequeue();
this.executeTask(task);
}
}
}
}
为什么需要模型级别隔离？因为同一个模型可能共享 GPU Buffer，并发写入会数据覆盖。
2. WebGPU 后端
WebGPU 是浏览器的新一代图形/计算 API，相比 WebGL：
- 原生支持计算着色器
- 更低的 CPU 开销
- 更精确的同步控制
核心是用 WGSL 写计算着色器：
@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
let idx = gid.x;
output[idx] = activation(input[idx] * weight[idx] + bias[idx]);
}
当然也做了降级：WebGPU → WebNN → WASM。
3. RAII 内存管理
借鉴 C++ 的 RAII 思想，实现了 MemoryScope：
const result = await withMemoryScope(async (scope) => {
// 在作用域内创建的 tensor 会自动追踪
const a = scope.track(tensor([1, 2, 3]));
const b = scope.track(tensor([4, 5, 6]));
// 处理...
const output = matmul(a, b);
// 标记要保留的结果
return scope.keep(output);
});
// 离开作用域，a 和 b 自动释放
还做了泄漏检测：记录每个资源的创建时间和调用栈，超时未释放就报警。
4. 分片下载 + 断点续传
大模型动辄几百 MB，传统下载体验很差。我实现了：
- 分片下载：用 HTTP Range 请求并行下载多个 chunk
- 断点续传：下载进度存 IndexedDB，刷新页面继续
- 智能缓存：下载完的模型缓存到 IndexedDB，下次秒加载
// 并行下载
const chunks = await Promise.all([
downloadChunk(url, 0, 5MB),
downloadChunk(url, 5MB, 10MB),
downloadChunk(url, 10MB, 15MB),
downloadChunk(url, 15MB, 20MB),
]);
5. 流式文本生成
支持 GPT 风格的 token-by-token 输出：
for await (const event of generator.stream('Once upon a time')) {
process.stdout.write(event.token); // 逐字输出
}
实现用了 AsyncGenerator，配合 top-k/top-p 采样。
学到了什么
做这个项目让我深入理解了：
- 浏览器并发模型 - 事件循环、微任务队列、Worker 通信
- WebGPU 编程 - 计算着色器、Buffer 管理、同步
- 内存管理 - JS 的 GC 机制、手动资源管理的必要性
- 大文件处理 - 分片、断点续传、IndexedDB 的坑
如何使用
安装：
npm install edgeflow
使用：
import { pipeline } from 'edgeflow';
// 情感分析
const classifier = await pipeline('text-classification');
const result = await classifier.run('I love this!');
// 并发执行
const [r1, r2] = await Promise.all([
classifier.run('Great!'),
extractor.run('Hello')
]);
// 流式生成
for await (const {token} of generator.stream('Hello')) {
console.log(token);
}
最后
这是我第一个认真做的开源项目，花了很多心血。如果觉得有用：
- ⭐ 给个 Star 支持一下
- 🐛 遇到问题欢迎提 Issue
- 🔀 欢迎 PR 贡献代码
GitHub: github.com/s-zx/edgeFl…
有问题欢迎评论区交流！
