1. 为什么非要把 AI 检索塞进浏览器里
最近几年端侧 AI 这个概念算是彻底火了,从手机上的 NPU 到浏览器里的 WebGL,模型推理正在从服务器大规模迁移到用户设备上。我这次要分享的项目,就是在这个大背景下做的一个尝试:完全在浏览器端实现一套 1024 维视觉向量特征检索系统。简单说,用户在网页里上传一张图片,系统在本地提取特征向量,然后和一个预先构建好的向量库做相似度匹配,返回最接近的若干结果。整个过程不经过任何服务器,模型、图片、向量、索引全部停留在用户自己的设备上。
这个需求的场景其实很常见。比如电商网站的商品相似推荐、照片管理工具的自动归档、企业内部的知识库图片搜索。传统做法一般是把图片上传到云端,用 Python 那套技术栈跑模型,再用 Milvus、Faiss 这类向量数据库做召回。但问题也很明显:第一,每个用户的特征提取、向量存储、检索计算都要消耗云端资源,用户量一上来账单就非常可观;第二,用户图片这类敏感性数据经过网络传输和云端存储,隐私上总归是个隐患,有些行业(比如医疗影像、法务材料)压根不允许数据出内网;第三,网络往返带来的延迟在一些需要实时反馈的场景里也不够理想。
我当时接到这个需求时,产品那边给了两个硬性要求:图片和检索行为必须全程留在本地,而且新增一个用户不能给服务器带来任何额外计算开销。说白了就是零云端成本加 100% 隐私安全。那方案就非常清晰了——直接在浏览器里做推理和检索。适合参考这个方案的人我觉得有这么几类:被云成本压得喘不过气的前端团队,正在做隐私敏感型应用的全栈工程师,还有那些想在浏览器里玩模型推理但不知道从哪下手的 AI 爱好者。看完这篇文章,你至少能搞清楚整个技术闭环怎么串起来,以及实际落地时会踩到哪些坑。
2. 整体设计与技术选型:为什么是 TensorFlow.js 加 Web Worker
2.1 端侧推理方案摆在一起比一比
把 AI 推理搬到浏览器里,现在能选的路子其实不少。我简单梳理一下当时对比过的几个方案。首先是 TensorFlow.js,它提供了 WebGL、WebAssembly、WebGPU 三套后端,API 风格接近 Keras,前端开发者上手门槛低,而且模型转换工具链非常成熟,不管是 TensorFlow SavedModel 还是 Keras H5,都能转成浏览器可加载的格式。其次是 ONNX Runtime Web,如果你手里的模型是 PyTorch 导出的 ONNX 格式,这套运行时也很香,推理性能在某些模型上甚至比 TensorFlow.js 还好,但它对开发者的要求更高,需要你自己做更多底层的内存管理和张量布局优化工作。还有纯 WebGPU 手写推理的方案,性能天花板最高,但开发成本完全是另一个量级,适合那种需要极致性能且人力充裕的团队。
我最后选 TensorFlow.js,核心原因倒不是因为它的推理速度最快,而是生态最省心。这个项目里我需要的是一个开箱即用的图像特征提取模型,TensorFlow.js 官方模型库里有现成的 MobileNet 实现,加载一行代码,预处理逻辑也封装好了。考虑到团队里成员的主业是前端,让他们去手写 WebGPU shader 显然不现实。另外 TensorFlow.js 的 WebGL 后端能在绝大多数浏览器上直接跑,不需要用户开启任何 flag,这对产品落地来说太重要了。
2.2 1024 维特征向量是怎么来的
项目标题里提到的 1024 维向量,来源于模型结构里最后一个池化层的输出。我用的是 MobileNet V2 作为特征提取器,输入图片会被 Resize 到 224×224 这种标准尺寸,经过模型一系列深度可分离卷积之后,在倒数第二层会得到一个 1×1×1280 的特征图。这里其实有个细节,如果你直接拿 MobileNet 官方模型的分类输出(1000 个类别的概率分布),那得到的是 1000 维向量,而且语义信息高度集中在分类任务上;但如果你把网络截断在最后一个全局平均池化层,那么得到的是 1280 维的图像语义特征。我项目里最终用 1024 维,是因为换用了 MobileNet V3 的一个变体,它在倒数第二层正好输出 1024 维的 embedding,兼顾了模型体积和特征表达力。
提取特征向量之后还有一步很重要:对向量做 L2 归一化。这步操作的意义在于,后续检索用的余弦相似度计算,在向量归一化之后可以简化为点积运算,省去每次计算模长的开销。更关键的是,归一化能削弱图片整体亮度、对比度对特征的影响,让检索结果更关注语义内容。我当时实测下来,同一个物体在不同光照条件下拍两张照片,归一化之后特征的余弦相似度从 0.82 提升到了 0.9 以上,效果还是很明显的。
2.3 Web Worker 其实是整个应用能用的关键
很多人一开始不理解,TensorFlow.js 不是已经支持异步了吗,为什么还要引入 Web Worker。这个认知有个误区:TensorFlow.js 的异步 API 只是说 JavaScript 线程不会被阻塞,但实际上模型推理这种计算密集型任务运行时,会占用主线程的 CPU 时间片和 GPU 资源。我做过一个简单测试:在主线程上用 mobileNet 跑一次推理,页面如果同时有滚动动画或者视频播放,帧率直接掉到个位数,整个页面跟卡死了一样。原因是浏览器的主线程要同时处理 DOM 渲染、事件响应和 JavaScript 执行,任何一个重任务挤进来都会引起明显的交互延迟。
把推理逻辑放进 Web Worker 之后,情况就完全不一样了。Worker 运行在独立的线程上,有自己的全局上下文,不碰 DOM 也不影响主线程的事件循环。不过这里有个容易踩的坑:WebGL 上下文在主线程和 Worker 之间是不能共享的,所以在 Worker 里跑模型推理,实际上会创建第二个 WebGL 上下文。这是可行的,但你要在 Worker 初始化时单独加载模型,并且和主线程之间通过 postMessage 通信。我们的架构里,主线程负责图片文件读取和结果渲染,Worker 负责模型加载、特征提取、向量检索这三件核心计算工作。实践中这样拆分后,页面主线程几乎感知不到背后在跑模型,流畅度问题直接消失。
2.4 为什么不用矢量数据库而自己写检索
既然标题里说到了向量检索,可能有人会问,为什么不直接在浏览器里嵌一个 Faiss 的 WASM 编译版本。这个方向确实可行,GitHub 上也有类似项目。但对我来说,Faiss 是面向服务端设计的库,它的优势在于大规模索引构建、GPU 加速、IVF 这种高级索引结构。而浏览器端的场景,向量库规模一般不会特别大,个人相册几千张图、企业知识库几万张图,这个量级用暴力线性扫描(Brute Force)就足够了。1024 维向量的余弦相似度计算,单次点积大约 1024 次浮点运算,加上归一化预处理,网上文章里最常见的估算是一万条向量全量扫描耗时约几十毫秒,我实测在 M1 芯片的浏览器里处理两万条索引也就 120 毫秒左右。这个性能完全够交互使用,而且代码简单、没有额外依赖、不容易出错。
高维向量检索有个著名的现象叫维度灾难,意思是当维度升高到一定程度,各种空间索引结构(比如 KD 树)的查询效率会退化到接近线性扫描,甚至因为回溯开销变得更慢。1024 维绝对属于高维范畴了,所以在这个项目里,简单的线性扫描反而成了最优解。等到以后向量库规模真的到了十万、百万级别,再考虑引入更复杂的索引也不迟。
3. 核心实现细节与实操过程
3.1 图片特征提取的完整链路
整个特征提取链路可以拆成四步:图片解码、预处理、模型推理、后处理。下面是 Worker 内部一个简化的核心函数:
// feature-extract.js(运行在 Worker 内) import * as tf from '@tensorflow/tfjs'; let model; // 模型加载:使用 MobileNet 并截断到 embedding 层 async function loadModel() { model = await tf.loadLayersModel('/models/mobilenetv3/model.json'); // 这里的关键点:截取倒数第二层输出 const embeddingLayer = model.getLayer('embedding_output'); model = tf.model({ inputs: model.inputs, outputs: embeddingLayer.output }); } async function extractFeature(imageBitmap) { return tf.tidy(() => { // imageBitmap -> tensor,顺便做归一化 let tensor = tf.browser.fromPixels(imageBitmap); tensor = tf.image.resizeBilinear(tensor, [224, 224]); tensor = tf.cast(tensor, 'float32').div(127.5).sub(1); // 增加 batch 维度:从 [224,224,3] -> [1,224,224,3] tensor = tensor.expandDims(0); // 推理得到 embedding const embedding = model.predict(tensor); // L2 归一化,让检索阶段可以直接用点积 const normalized = tf.div(embedding, tf.norm(embedding, 'euclidean')); // 转成普通 Float32Array,用于后续存储和比较 return normalized.dataSync(); }); }这段代码有几个细节值得展开。首先,tf.browser.fromPixels接收的是 ImageBitmap、HTMLImageElement 或者像素数组,在 Worker 里不能直接访问 DOM,所以主线程需要先把图片文件转成 ImageBitmap,再通过 postMessage 传进 Worker。其次,tf.tidy()是个好东西,它能自动回收函数内部创建的中间张量内存,避免 WebGL 上下文中的显存泄漏。如果你不用 tidy,每个中间张量都得手动调dispose(),稍有不慎浏览器标签页就会因为显存暴涨被系统干掉。
针对模型截断这一步我再补充一下,tf.model可以基于已有模型创建子模型,只保留输入层到指定层的计算路径。但前提是你转换模型时,各个层要有清晰的名称。如果模型是转换工具自动生成的,建议先打印model.summary()看一眼层名称,确认 embedding 层叫什么,再写截断逻辑。我当时一开始用model.layers[model.layers.length - 3]这种索引方式取层,结果模型一更新就报错,改成按名称获取之后才稳定下来。
3.2 向量检索核心算法:余弦相似度的正确打开方式
检索阶段的任务是给定一个查询向量,从索引库中找出最相似的 Top-K 条记录。由于前面已经做了 L2 归一化,余弦相似度退化成向量点积,实现极其简洁:
// search.js(运行在 Worker 内) // indexData: [{ id, vector: Float32Array, meta }] function search(queryVector, indexData, topK = 10) { const scores = new Array(indexData.length); for (let i = 0; i < indexData.length; i++) { scores[i] = { id: indexData[i].id, // 简单的点积 = 余弦相似度(因为两边都已归一化) score: dotProduct(queryVector, indexData[i].vector), meta: indexData[i].meta }; } scores.sort((a, b) => b.score - a.score); return scores.slice(0, topK); } function dotProduct(a, b) { let sum = 0; for (let i = 0; i < a.length; i++) { sum += a[i] * b[i]; } return sum; }这段代码的性能瓶颈在 JavaScript 循环上,尤其是向量库达到几万条时,点积运算是纯浮点计算,JS 引擎编译优化后其实也还行,但我们可以再优化一把。第一,把向量库存成一个大号 Float32Array 而不是数组里面套数组,这样内存连续,CPU 缓存命中率高;第二,用for循环而不是forEach或map,避免函数调用开销;第三,如果浏览器支持 SIMD(WebAssembly 的 SIMD 指令),可以把点积挪进 WASM 里写,实测可以再快 2 到 4 倍。不过我在这个项目里先用纯 JS 顶着,因为两万条数据毫秒级的检索时间已经够用了。
还有一个容易被忽略的点:检索阈值的设定。余弦相似度的范围是 -1 到 1,但实际图片特征向量的相似度分布很少拉满。我测试下来,同一场景不同角度的图片,相似度在 0.7 到 0.85 之间;完全不同的内容大约在 0.1 到 0.4 之间。所以检索时一般设一个 0.6 左右的阈值,低于这个阈值的记录直接过滤,避免返回一堆毫不相关的噪声结果。这个阈值肯定不是固定的,最好做成可配置参数,用户反馈精度不够就调高,召回不够就调低。
3.3 向量数据持久化:IndexedDB 里存二进制
索引数据不能每次打开页面都重新计算,那太浪费了。对个人相册这种场景,理想状态是第一次导入图片时提取特征,之后直接复用。我选择用 IndexedDB 做持久化存储,把图片原始文件(Blob)和对应的特征向量分开存放。
这中间有一个重要的经验:向量不要存成 JSON 数组。1024 个浮点数如果用 JSON 存,会写成类似[0.00123124, 0.98231232, ...]的文本,占空间大且解析慢。正确做法是直接把 Float32Array 的底层 ArrayBuffer 存进 IndexedDB,读取的时候用new Float32Array(buffer)包一层就能直接用。这样存储体积是 JSON 的四分之一(二进制 4 字节一个 float,JSON 文本平均至少 15 字节一个数),解析时间几乎为零。
IndexedDB 的写入和读取都是异步的,正好适配 Worker 里的异步模型。初始化 Worker 时先把索引数据整体加载到内存,做成 Float32Array 的矩阵。内存占用方面,一条 1024 维向量是 4KB,一万条就是 40MB,这在现代设备上完全可以接受。如果你的向量库更大,就得做内存映射或者分页加载了,不过那是另一个复杂度等级的问题。
3.4 主线程与 Worker 之间的消息协议设计
主线程和 Worker 之间通过 postMessage 传递消息,消息数量不多的时候随便传都行,但一旦功能多了,还是要设计一套简单的协议,不然代码就成了大型 if-else 现场。我定义的协议大概是这样的结构:
// 主线程发往 Worker 的消息 const MessageType = { LOAD_MODEL: 'load_model', EXTRACT_AND_ADD: 'extract_and_add', // 提取特征并加入索引 SEARCH_BY_IMAGE: 'search_by_image', // 输入一张图,检索相似图 SEARCH_BY_VECTOR: 'search_by_vector', // 输入一个向量,检索相似向量 DELETE_BY_ID: 'delete_by_id', // 删除索引 REBUILD_INDEX: 'rebuild_index' // 重新构建索引 }; // Worker 返回给主线程的消息结构 // { // type: 'search_result', // requestId: 123, // result: [{ id, score, meta }] // }这里有个很实用的设计:每条请求带上requestId,Worker 处理完返回时把 requestId 带上。这样主线程里可以做请求和响应的匹配,避免多个并发请求回来之后乱套。比如用户快速点击了三次不同图片进行检索,Worker 处理顺序可能是乱序的,如果没有 requestId,主线程根本不知道哪条响应对应哪次请求。
消息体里传递图片数据时,尽量传 ImageBitmap 而不是 File 或者 Blob。ImageBitmap 是经过解码的位图数据,在 postMessage 时浏览器会做优化,不需要重新解码图片格式。如果你的图片是 JPEG 这种压缩格式,Worker 里拿 Blob 还得走一遍 createImageBitmap 解码流程,多了等待时间。主线程先把 File 变成 ImageBitmap,再传给 Worker,这条链路明显更快。
4. 从零搭建到能用的完整实操过程
4.1 项目初始化与依赖安装
项目我用 Vite 做构建工具,主要是看中它的开发服务器启动快,以及原生支持 Web Worker 模块化写法。Vite 环境下 Worker 可以直接用new Worker(new URL('./worker.js', import.meta.url), { type: 'module' })这种方式引入,开发体验比 Webpack 的 worker-loader 舒服不少。
依赖方面就一个@tensorflow/tfjs。模型文件我放在 public 目录下静态托管,用相对路径加载。如果你有现成的 TensorFlow SavedModel 模型,可以用官方的 tfjs-converter 转成浏览器格式:
pip install tensorflowjs tensorflowjs_converter \ --input_format=tf_saved_model \ --output_format=tfjs_graph_model \ /path/to/saved_model \ /path/to/web_model转换完成之后,目录里会有一个model.json和一组分片的bin文件。注意model.json里会记录权重文件的路径,如果路径里的 URL 前缀不对,浏览器会加载失败。我建议转换时就用相对路径的模型目录,别在 model.json 里写绝对路径,部署之后遇到跨域问题会很头疼。
4.2 Worker 内部的初始化流程
Worker 内部启动时要完成模型加载和索引加载两件事。模型加载是重量级操作,MobileNet V3 的模型文件大约 15MB 左右(分片权重文件),首次加载需要从网络拉取,后面可以用 Cache API 缓存到浏览器本地。索引加载是从 IndexedDB 里读出历史数据,构建内存中的向量矩阵。
我定义了一个启动时序:Worker 收到第一个消息前先把模型加载好,加载完成后主动向主线程发一条ready消息。主线程收到 ready 之前,会把所有请求放进队列缓存起来,等 Worker 就绪后统一发送。没有这个机制,刚刷新页面时用户立刻点检索,消息发过去时模型还没加载完,就会报错。其实这段逻辑也不复杂,就是多了一个状态管理的问题:主线程维护一个workerReady标志位,没就绪就先缓存请求。
模型加载的进度反馈也是个产品体验细节。tf.loadLayersModel支持传onProgress回调,能拿到已加载字节数和总字节数,可以用它来做一个简单的进度条。不然用户看到页面黑屏空白十几秒,不知道在干嘛,早就关页面了。
4.3 索引构建的批处理优化
首次使用时要批量导入图片,比如导入一千张图片到相册里。如果一张一张地提取特征再存储,每一张都要走一遍图片解码、模型推理、IndexedDB 写入的流程,总耗时可能超过两分钟,用户等不起。
我做了一个分批处理机制:每批次处理 8 张图片,用Promise.all并发执行特征提取。这么做既利用了现代 CPU 的多核能力,又不会因为并发数太大导致内存暴涨。每个批次完成后更新一次进度,并分批写入 IndexedDB。实践中,一千张图片在较新的 MacBook 上大概 40 秒到一分钟能完成,至少不会让用户怀疑脚本死掉了。
还有一个细节是处理图片前先做压缩。手机随便拍的照片都是 12MP 起步,如果原图直接送进模型,解码开销很大但特征质量和 224×224 输入差不多。我额外加了一步:超过 1024 像素边长的图片先画到小一点的 canvas 上,等比缩到最长边 1024 再提取特征。这个优化让单张图片的预处理时间降低了 60% 以上,对检索精度几乎没有可感知的影响。
4.4 多检索结果的可视化
检索结果的可视化放在主线程,Worker 只负责返回匹配项的 id、相似度分数和 meta 信息。meta 里我存储了图片名称、缩略图 Blob 的 IndexedDB 主键。拿到结果后,主线程从 IndexedDB 里按 id 取出缩略图,直接渲染到页面网格上。
相似度分数这一栏建议保留两位小数展示,别显示那种 0.8372641 的长串数字。用户对 0.83 和 0.84 的差别已经感知不明显了,更重要的是在结果卡片上标注相似度排名。有条件的还可以做一个差异叠加视图,把查询图片和结果图片重叠显示,用半透明效果让用户直观看到哪些区域匹配上了。这在文物修复、部件比对这类专业场景里格外有价值。
5. 常见问题与排查技巧实录
5.1 问题速查表
我把开发过程中遇到频率最高的几个问题整理成了表格,你在复现时如果撞上了可以直接按图索骥。
| 现象 | 根本原因 | 解决方案 |
|---|---|---|
| 模型加载永远卡在 99% | 权重分片文件跨域请求被拦截 | 模型文件放到同源目录,或配置 CORS 头 |
| Worker 内推理显存持续增长 | 缺少 tf.tidy() 或张量未释放 | 用 tf.tidy 包裹推理逻辑,手动 dispose 中间张量 |
| 移动端检索结果明显变差 | WebGL 浮点精度低,iOS 上部分机型为 fp16 | 切换 WASM 后端,或用 fp32 纹理扩展 |
| 首次加载模型白屏十几秒 | 模型文件未缓存到本地 | 用 Cache API 缓存 model.json 和权重,二次加载走缓存 |
| postMessage 传递大图片卡顿 | 传输了未压缩的 File 或超大地图 | 先转成 ImageBitmap 或 canvas 压缩后再传递 |
| IndexedDB 存储占用爆炸 | 把 Float32Array 转成 JSON 存储 | 直接存 ArrayBuffer 二进制,读取时再包 Float32Array |
| Worker 启动后 CPU 占用过高 | 模型加载和特征提取同时并发 | 分批处理,限制并发数,空闲时自动降频 |
5.2 移动端 WebGL 的精度陷阱
我在 Android 和 iOS 设备上测试时发现了一个有趣的现象:同一张图片,在桌面端 Chrome 和手机 Safari 上提取出来的特征向量竟然有细微差别,直接导致检索结果的排名顺序不完全一致。查了一圈资料,原因是移动端 GPU 的 WebGL 实现可能采用 fp16 半精度浮点纹理,导致中间层计算结果精度损失。这个问题在 MobileNet 这种容错性强的模型上影响不大,但如果换成更敏感的特征模型,结果可能就差之毫厘谬以千里。
解决办法是强制 TensorFlow.js 使用 WASM 后端,或者启用WEBGL_*相关的精度扩展。WASM 后端走 CPU 计算,虽然速度比 GPU 慢,但精度和桌面端能保持完全一致。如果检索结果的一致性是你的硬需求,建议在移动端做后端自动降级:检测到navigator.userAgent是移动设备就用 WASM,桌面端用 WebGL。这个方案在兼容性和性能之间算是比较好的平衡点。
5.3 索引更新时的并发控制
在线索引更新有多用户同时写入的场景时,Worker 内部要注意消息处理的串行化。我刚实现时很天真地写成了纯并发模型,dispatch 出去就直接返回,结果出现索引竞争:两个请求同时往索引数组中 push,导致检索时数组越界或者数据丢失。
后来我引入了一个简单的任务队列:所有写操作(添加、删除)串行执行,读操作(搜索)可以并发。每个写操作执行时,先把内存中的索引复制一份(写时复制),修改完成后再原子地替换引用。这样读者要么看到旧版本索引,要么看到新版本索引,永远不会看到写入一半的非法状态。这个模式在浏览器环境里没有锁可用,但借助 JavaScript 的单线程事件循环特性,只要保证引用替换是同步的,就能实现无锁读并发。
5.4 大规模索引的多 Worker 分片方案
当索引规模真的到了十万条以上,单 Worker 的线性扫描就会明显吃力。我的压测数据是五万条向量,单 Worker 检索耗时大约八百毫秒,对交互应用来说已经有点卡顿了。这时候可以把索引分片到多个 Worker,每个 Worker 加载一部分数据,搜索时主线程把查询向量广播给所有 Worker,各个 Worker 并行执行局部检索,最后把各自 Top-K 结果汇总到主线程再排序。
多 Worker 分片的内存开销是翻倍的,因为每个 Worker 都要复制一份自己的索引分片。这里有一个进阶技巧:用 SharedArrayBuffer 共享底层向量数据,多个 Worker 读同一块内存,不需要各自复制。但 SharedArrayBuffer 要求页面开启跨源隔离(COOP/COEP 响应头),部署在静态托管服务上时需要改服务器配置。这个方向我目前只在技术验证阶段试过,生产环境还没铺开,如果你正好要冲十万级以上的数据量,可以沿着这条路继续深挖。
我在实际开发过程中的一个体会是,这套端侧架构一旦跑通,维护成本比同等的云端服务低一个量级。模型更新只需要重新发布静态资源,不需要考虑灰度发布时的流量控制;存储扩容是用户本地磁盘的事,服务器这边完全无感。项目做完之后我其实最欣慰的倒不是省下了多少钱,而是客户对数据安全的担心彻底解除了,因为这个方案从物理层面就不存在数据离开设备的那条通道。最后再分享一个小技巧:调试 Worker 代码时,在 Chrome DevTools 的 Sources 面板里找到 Worker 的上下文,可以像调试主线程一样打断点、看变量,别在代码里塞一堆 console.log 之后再删了,那个效率实在太低了。