你有沒有遇到過這種情況客戶要一個“以圖搜圖”功能圖片又不想傳到服務(wù)器上要求在瀏覽器里直接完成而且速度還不能慢。傳統(tǒng)前端想都不敢想的任務(wù)現(xiàn)在靠瀏覽器端 AI 加上 WebGPU 加速已經(jīng)能跑得很舒服了。我這次做的項目就是讓瀏覽器加載一個 MobileNet 特征提取模型把用戶上傳的圖片變成 1024 維向量再拿這個向量去全本地索引里搜相似圖整個鏈路不依賴 GPU 服務(wù)器也不把圖片發(fā)出去。這套方案的核心技術(shù)點有三個MobileNet 這種輕量 CNN 為什么適合在瀏覽器端跑WebGPU 和傳統(tǒng) WebGL 推理的差異到底在哪里以及 1024 維特征向量怎么支撐起一套可用的以圖搜圖鏈路。適合看這篇文章的人我大概歸了三類想做前端 AI 應(yīng)用但卡在性能和兼容性上的開發(fā)者正在做素材管理、相冊去重、版權(quán)比對這類圖片檢索需求的工程師還有想低成本給自己的網(wǎng)頁應(yīng)用加“相似圖片”搜索能力的產(chǎn)品和技術(shù)決策者。文章里的代碼基于 onnxruntime-web 實現(xiàn)模型用 MobileNetV2 結(jié)構(gòu)推理走 WebGPU向量檢索在前端用純 JavaScript 完成。我會把從模型導出、環(huán)境準備、代碼拆解到踩坑記錄全流程講一遍你跟著操作就能在自己項目里復(fù)現(xiàn)。1. 為什么非要在瀏覽器里跑特征提取1.1 用戶要的“不傳圖”實際上要的是隱私和速度先說個挺有意思的現(xiàn)象。做后端推理的兄弟可能覺得客戶端傳圖、服務(wù)器返回結(jié)果不是挺成熟的方案嗎確實成熟但落地的時候有三個問題很難繞開。第一是隱私。相冊類應(yīng)用、企業(yè)內(nèi)部素材庫、醫(yī)療影像系統(tǒng)的用戶對圖片數(shù)據(jù)極其敏感圖片一旦上行合規(guī)、加密、審計的流程就全來了。很多項目根本不是技術(shù)做不到是“圖片能不能出域”這一條直接被客戶否掉了。第二是延遲。圖像上傳要看帶寬服務(wù)器排隊要看負載。哪怕后端推理只要 50 毫秒一張兩兆的圖傳上去再加上網(wǎng)絡(luò)抖動用戶感知到的就是“轉(zhuǎn)圈”。而瀏覽器端直接處理圖片從本地文件到特征向量幾百毫秒內(nèi)完成這種反饋是實打?qū)嵉?。第三是成本。做圖片特征抽取要吃 GPU 資源如果只是偶爾查一次專門起一個 GPU 服務(wù)就是浪費。把計算挪到用戶設(shè)備上服務(wù)器只負責靜態(tài)資源分發(fā)成本結(jié)構(gòu)完全不一樣。所以“瀏覽器端跑 AI”聽起來像炫技實際上是在隱私、速度、成本三個維度上都更優(yōu)的選擇。尤其是圖庫規(guī)模不大比如幾千到幾萬張、查詢頻率不高的場景純前端方案是性價比最高的一條路。1.2 WebGPU 相比 WebGL 到底強在哪可能有人問瀏覽器端跑神經(jīng)網(wǎng)絡(luò)又不是今天才有WebGL 時代就有人用 TensorFlow.js 跑過模型了。為什么現(xiàn)在要提 WebGPU因為這個差距實在太大了。WebGL 本質(zhì)上是為二維和三維圖形渲染設(shè)計的 API想在它上面跑神經(jīng)網(wǎng)絡(luò)得把卷積、矩陣乘法這些算子偽裝成片段著色器把數(shù)據(jù)塞進紋理再把計算結(jié)果從紋理里讀回來。這個過程中紋理讀寫的限制很多內(nèi)存分配和算子控制也很粗糙遇到復(fù)雜的網(wǎng)絡(luò)結(jié)構(gòu)性能損耗非常明顯。WebGPU 不一樣它引入了 compute shader 的概念能夠直接操作 GPU buffer做通用并行計算。這意味著卷積、矩陣乘法這類重計算算子可以用更接近底層的方式組織數(shù)據(jù)不必反復(fù)在紋理和 CPU 之間倒騰。對神經(jīng)網(wǎng)絡(luò)推理來說WebGPU 還有一個天然優(yōu)勢支持更靈活的 tensor 布局比如直接用 NCHW 布局減少推理框架內(nèi)部的 transpose 開銷。我用同一個 MobileNet ONNX 模型做過對比在同樣的機器上WebGL 后端跑一次推理大約 300 到 500 毫秒WebGPU 后端降到 30 到 60 毫秒差距接近一個數(shù)量級。這種差距決定了“在瀏覽器里做實時圖片檢索”是可行的還是只能停留在玩具階段。1.3 選型MobileNet ONNX Runtime Web 的組合怎么定下來的模型選 MobileNet理由很直接。MobileNet 的核心設(shè)計是深度可分離卷積把標準卷積拆成逐通道卷積和逐點卷積兩步參數(shù)量和計算量大幅下降很適合瀏覽器這種算力受限的環(huán)境。MobileNetV2 在 ImageNet 上的精度雖然不如 ResNet 這類大模型但作為特征提取器它的瓶頸層輸出已經(jīng)能很好地表達圖像語義足夠支撐相似圖片檢索。推理引擎我選了 ONNX Runtime Web。這里要說說為什么不用 TensorFlow.js。TensorFlow.js 對 TensorFlow 模型生態(tài)的支持確實好但我日常更多是 PyTorch 訓練模型ONNX 作為中間格式更通用。onnxruntime-web 對 ONNX 算子的覆蓋度一直在提升而且從 1.17 版本開始加入 WebGPU 執(zhí)行后端走的是官方維護的路線。特征維度定在 1024不是拍腦袋。維度太低相似圖片容易誤判維度太高索引和檢索的開銷明顯增加。人臉識別領(lǐng)域常見的特征維度是 512 或者 1024這個區(qū)間在語義區(qū)分度和計算量之間比較平衡。MobileNetV2 原版瓶頸層輸出是 1280 維我在項目里對模型做了一點調(diào)整輸出層改成 1024 維的 embedding 頭后面會詳細說導出方式。2. 動手前的準備環(huán)境、模型導出與工程骨架2.1 瀏覽器端 WebGPU 能不能用的自檢方法寫代碼之前先確認目標環(huán)境下 WebGPU 到底可不可用。WebGPU 需要瀏覽器和硬件驅(qū)動雙重支持不是裝了新瀏覽器就一定行。最直接的檢測方法是看navigator.gpu是否存在。Chromium 內(nèi)核的瀏覽器從 Chrome 113 開始默認支持 WebGPUEdge 113 也同步支持。Firefox 目前是部分支持Safari 在較新版本里也開始跟進。開發(fā)階段建議優(yōu)先用 Chrome 或 Edge 調(diào)試。還有一個隱藏條件WebGPU 要求安全上下文。也就是說頁面必須在 HTTPS 環(huán)境下打開或者本地 localhost 調(diào)試普通的 HTTP 頁面是拿不到 GPU 權(quán)限的。我寫過一個簡單的檢查函數(shù)頁面加載時直接跑一遍function checkWebGPU() { if (gpu in navigator) { console.log(WebGPU 可用GPU 設(shè)備:, navigator.gpu); return true; } console.warn(當前瀏覽器不支持 WebGPU請使用 Chrome 113 或 Edge 113); return false; }如果navigator.gpu存在但推理速度異常慢建議打開chrome://gpu頁面看一下 WebGPU 相關(guān)的狀態(tài)和驅(qū)動告警。有時候是顯卡驅(qū)動版本太舊WASM 和 WebGPU 本身沒問題但硬件加速沒跑起來。2.2 把 MobileNet 導出成 1024 維特征模型這一步把 PyTorch 模型轉(zhuǎn)成 ONNX 格式。我用的 MobileNetV2 來自 torchvision默認的分類頭輸出 1000 類現(xiàn)在要把它替換成 1024 維的特征輸出。注意一個細節(jié)MobileNetV2 在全局池化之后classifier包含兩層第一層是 Dropout第二層是Linear(last_channel, 1000)其中l(wèi)ast_channel是 1280。要改成 1024 維特征直接替換最后的 Linear 層即可。import torch import torchvision import torch.nn as nn model torchvision.models.mobilenet_v2(weightstorchvision.models.MobileNet_V2_Weights.DEFAULT) model.classifier[1] nn.Linear(model.last_channel, 1024) model.eval() dummy_input torch.randn(1, 3, 224, 224) torch.onnx.export( model, dummy_input, mobilenet_v2_feat.onnx, input_names[input], output_names[feature], opset_version17 ) print(導出完成輸出維度應(yīng)在 ONNX 中為 [1, 1024])這里要提醒一句隨機初始化的 Linear 頭在沒經(jīng)過微調(diào)時特征質(zhì)量不一定比原版 1280 維好。我在實際項目里用的是在業(yè)務(wù)數(shù)據(jù)集上用度量學習微調(diào)過的 1024 維 head。如果你只是驗證流程可以暫時用隨機初始化的 head相似圖片檢索的結(jié)果會有一點下降但流程能跑通。另一個方案是直接用原版 1280 維輸出在前端用固定投影矩陣降維到 1024這個后面再聊。MobileNet 的輸入預(yù)處理要和訓練時保持一致resize 到 224x224像素歸一化到 0 到 1然后按 ImageNet 的 mean 和 std 做標準化。導出 ONNX 時沒有把歸一化層放到模型里所以這些步驟要在前端預(yù)處理完成。2.3 工程結(jié)構(gòu)、依賴與 wasm 文件放置項目使用 npm 安裝 onnxruntime-web 的穩(wěn)定版本。注意 WebGPU 后端還處于快速迭代期API 有變化建議鎖定一個大版本不要每次盲目升級。npm install onnxruntime-web1.19.2onnxruntime-web 底層是 WASM瀏覽器需要加載對應(yīng)的 wasm 文件。最簡單的做法是把node_modules/onnxruntime-web/dist/下的 wasm 文件復(fù)制到項目的public/wasm/目錄然后用ort.env.wasm.wasmPaths指定路徑。項目目錄結(jié)構(gòu)參考project/ ├── public/ │ ├── models/ │ │ └── mobilenet_v2_feat.onnx │ └── wasm/ │ ├── ort-wasm-simd-threaded.wasm │ └── ... ├── src/ │ ├── feature-extractor.js │ ├── image-search.js │ └── main.js └── index.htmlfeature-extractor.js負責加載模型、圖像預(yù)處理和推理image-search.js負責向量歸一化、相似度計算和圖庫管理。這樣拆開之后模型相關(guān)的邏輯和檢索邏輯可以獨立測試。初始化代碼是這樣的import * as ort from onnxruntime-web/webgpu; ort.env.wasm.wasmPaths /wasm/; const session await ort.InferenceSession.create(/models/mobilenet_v2_feat.onnx, { executionProviders: [webgpu], graphOptimizationLevel: all });這里把executionProviders指定為[webgpu]onnxruntime-web 會用 WebGPU 后端執(zhí)行算子。如果某些算子不支持 WebGPU會自動回退到 CPU。后續(xù)調(diào)試的時候可以打開日志確認算子有沒有發(fā)生回退。3. 核心代碼拆解從圖片到 1024 維向量3.1 圖像預(yù)處理resize、歸一化、布局轉(zhuǎn)換預(yù)處理是整個流程里最容易寫錯的地方因為模型對輸入布局極其敏感。ONNX 模型的輸入一般是 NCHW 格式也就是[batch, channel, height, width]在 JavaScript 里就是一個Float32Array數(shù)據(jù)順序是先放所有像素的 R 通道再放 G 通道最后放 B 通道。下面這段代碼把圖片畫到 canvas 上縮放到 224x224再按 ImageNet mean/std 做標準化輸出ort.Tensorasync function imageToTensor(imageSource) { const width 224; const height 224; const canvas document.createElement(canvas); canvas.width width; canvas.height height; const ctx canvas.getContext(2d, { willReadFrequently: true }); ctx.drawImage(imageSource, 0, 0, width, height); const imageData ctx.getImageData(0, 0, width, height); const data imageData.data; const float32Data new Float32Array(3 * width * height); const mean [0.485, 0.456, 0.406]; const std [0.229, 0.224, 0.225]; for (let i 0; i width * height; i) { const r data[i * 4] / 255.0; const g data[i * 4 1] / 255.0; const b data[i * 4 2] / 255.0; float32Data[i] (r - mean[0]) / std[0]; float32Data[width * height i] (g - mean[1]) / std[1]; float32Data[2 * width * height i] (b - mean[2]) / std[2]; } return new ort.Tensor(float32, float32Data, [1, 3, height, width]); }有幾個細節(jié)要提一下。getContext(2d, { willReadFrequently: true })這個參數(shù)在頻繁讀取像素時能避免 canvas 使用 GPU 加速導致的問題。如果圖片帶透明通道data[i * 4 3]是 alpha 值這里直接忽略因為 MobileNet 訓練時用的是 RGB 三通道圖像。如果你用的是createImageBitmap(file)讀取本地圖片得到的ImageBitmap對象可以直接傳給drawImage性能比先把文件轉(zhuǎn)成img標簽再畫要好尤其在批量處理場景下明顯。3.2 用 WebGPU 執(zhí)行推理并拿到特征向量預(yù)處理完之后推理代碼其實只有三行但背后發(fā)生的事情很多const output await session.run({ input: tensor }); const feature outputs[feature].data;session.run會走 WebGPU 計算管線輸入 tensor 上傳到 GPU buffer算子按圖執(zhí)行最后結(jié)果從 GPU buffer 讀回 CPU。第一次調(diào)用時WebGPU 需要完成管線編譯會有明顯的等待所以正常做法是在頁面加載后先“預(yù)熱”一次后面會專門講。拿到outputs[feature].data之后它是一個類型化數(shù)組長度理論上應(yīng)該是 1024。但 ONNX 輸出可能帶 batch 維度最好先確認一下形狀const shape outputs[feature].dims; // 例如 [1, 1024] const feature new Float32Array(outputs[feature].data); console.log(特征維度:, feature.length);這里用new Float32Array(...)復(fù)制一份很有必要。output.data可能引用的是 onnxruntime 內(nèi)部緩沖區(qū)后面再次推理時會被覆蓋復(fù)制出來可以避免臟數(shù)據(jù)。完整的抽取封裝大概長這樣let session null; async function initFeatureExtractor(modelPath) { if (session) return session; session await ort.InferenceSession.create(modelPath, { executionProviders: [webgpu], graphOptimizationLevel: all }); await warmup(); return session; } async function extractFeature(imageSource) { const tensor await imageToTensor(imageSource); const outputs await session.run({ input: tensor }); const feature new Float32Array(outputs[feature].data); return normalize(feature); }warmup就是拿一張全黑或全白的 224x224 圖先跑一次推理把 WebGPU 管線編譯和算子調(diào)度預(yù)熱完成這樣用戶真正查詢時就不會卡在那幾秒的編譯等待里。3.3 把特征向量喂給檢索模塊歸一化與余弦相似度特征向量做檢索之前先做 L2 歸一化。歸一化之后兩個向量的余弦相似度等于它們的內(nèi)積計算量小很多而且這種查詢方式在索引層面也更友好。export function normalize(vector) { let sum 0; for (let i 0; i vector.length; i) { sum vector[i] * vector[i]; } const norm Math.sqrt(sum) || 1; const normalized new Float32Array(vector.length); for (let i 0; i vector.length; i) { normalized[i] vector[i] / norm; } return normalized; } export function cosineSimilarity(a, b) { let dot 0; for (let i 0; i a.length; i) { dot a[i] * b[i]; } return dot; }注意如果a和b都已經(jīng)歸一化cosineSimilarity返回的就是余弦相似度值域在 -1 到 1 之間。以圖搜圖場景里返回結(jié)果一般按相似度從高到低排序。我在項目里會額外保存每個特征的平方和用來排查向量是否因為某種原因變成全零向量全零向量直接跳過不參與排序。4. 圖片庫構(gòu)建與檢索鏈路打通4.1 離線入庫批量抽取圖片特征以圖搜圖不止需要“查詢端”能抽特征圖片庫里的每張圖也得有特征。項目初期圖庫不大可以直接在瀏覽器里批量入庫把圖片庫的圖片逐張丟給extractFeature生成特征后存起來。批量入庫比單張查詢更容易做性能優(yōu)化。MobileNet 輸入是動態(tài) batch 維度的我可以把多張圖拼成一個[N, 3, 224, 224]的輸入張量一次推理同時輸出 N 個 1024 維向量。這個優(yōu)化在 WebGPU 后端下收益很大GPU 并行能力被充分用起來。批量特征抽取的示意代碼async function extractFeaturesInBatch(imageSources) { const tensor await batchImagesToTensor(imageSources); const outputs await session.run({ input: tensor }); const data outputs[feature].data; const shape outputs[feature].dims; const [batch, dim] shape; const features []; for (let i 0; i batch; i) { const vec new Float32Array(dim); vec.set(data.subarray(i * dim, (i 1) * dim)); features.push(normalize(vec)); } return features; }batchImagesToTensor就是把預(yù)處理的循環(huán)從單張變成多張構(gòu)造Float32Array(batch * 3 * 224 * 224)。我建議入庫過程放在 Web Worker 里跑避免 UI 卡死這個后面講工程化時展開。4.2 特征存儲IndexedDB 與二進制編碼入庫生成的特征需要持久化否則每次刷新頁面都要重新抽一遍特征。瀏覽器端的自然選擇是 IndexedDB它支持結(jié)構(gòu)化克隆Float32Array可以直接保存。我的存儲結(jié)構(gòu)很簡單每條記錄一個 id、一個Float32Array特征向量、還有圖片的縮略圖 Blob 和元信息。查詢時把圖庫全部向量取出來在內(nèi)存里做相似度計算。function saveFeatureToDB(id, vector, thumbnailBlob) { return new Promise((resolve, reject) { const request indexedDB.open(image-search-db, 1); request.onupgradeneeded () { const db request.result; if (!db.objectStoreNames.contains(features)) { db.createObjectStore(features, { keyPath: id }); } }; request.onsuccess () { const db request.result; const tx db.transaction(features, readwrite); tx.objectStore(features).put({ id: id, vec: vector, thumb: thumbnailBlob, createdAt: Date.now() }); tx.oncomplete () { resolve(); }; }; }); }要提醒一句IndexedDB 的結(jié)構(gòu)化克隆在保存Float32Array時是安全的但所有讀出來的Float32Array都是新的副本不會和內(nèi)存里的引用沖突這一點用起來很省心。如果你的項目圖庫超過十萬張1024 維 float32 的向量就是 400 多兆純前端暴力檢索的內(nèi)存壓力會非常大必須引入索引剪枝策略這個放到最后一節(jié)講。4.3 在線查詢一次完整的上傳檢索流程在線查詢的完整鏈路是用戶選擇圖片 -createImageBitmap讀取 -extractFeature抽特征 - 從 IndexedDB 讀取庫向量 - 逐個計算余弦相似度 - 排序取 TopK - 展示結(jié)果。檢索函數(shù)核心代碼async function searchByImage(imageSource, db, topK 10) { const queryVec await extractFeature(imageSource); const allFeatures await loadAllFeatures(db); const results []; for (const item of allFeatures) { const score cosineSimilarity(queryVec, item.vec); if (Number.isFinite(score)) { results.push({ id: item.id, score, thumb: item.thumb }); } } results.sort((a, b) b.score - a.score); return results.slice(0, topK); }細節(jié)上有一個性能優(yōu)化點查詢向量在循環(huán)外只歸一化一次庫向量在入庫時也已經(jīng)歸一化所以內(nèi)循環(huán)就是純乘加運算。對于一萬張圖就是一萬次 1024 維的內(nèi)積主線程上大概 10 到 20 毫秒體驗完全能接受。如果查詢過程中發(fā)現(xiàn)結(jié)果明顯不對先別急著改模型檢查一下查詢向量和庫向量是不是都做了歸一化這是我踩過最多的坑。5. WebGPU 推理的坑我替你踩了一遍5.1 打開頁面發(fā)現(xiàn) navigator.gpu 是 undefined這是最常見的兼容性問題。用戶瀏覽器版本太老或者頁面不在 HTTPS/localhost 環(huán)境下都會導致 WebGPU API 不可用。我的處理方式是在應(yīng)用入口做雙保險先檢查navigator.gpu不可用時給出提示同時給用戶一個回退按鈕用 CPU 執(zhí)行器跑同一個模型。const useWebGPU gpu in navigator; const executionProviders useWebGPU ? [webgpu] : [cpu];不過要把預(yù)期管理好CPU 后端跑 MobileNet一次推理大概在 200 到 400 毫秒小圖庫勉強能接受大圖庫就差很多。項目上線前一定提前確認目標用戶群用的瀏覽器版本不要指望所有用戶都有最新版 Chrome。5.2 首次推理慢得像假死預(yù)熱能解決多少WebGPU 的“首次運行慢”是新手遇到最多的現(xiàn)象明明同一個模型第二次跑快得很第一次卻可能要等好幾秒。原因是 WebGPU 后端需要把計算管線編譯成 GPU 指令這個編譯過程在第一次用一組算子形狀時觸發(fā)。方案就是預(yù)熱。session 創(chuàng)建完成后立即用一張和真實輸入尺寸完全一樣的假圖跑一次推理。這里強調(diào)“尺寸完全一樣”因為 WebGPU 管線編譯是按輸入形狀緩存的你用 224x224 預(yù)熱之后如果突然來一張別的尺寸輸入會再次觸發(fā)編譯。因此我在項目里把所有輸入都固定 resize 到 224x224避免形狀變化。預(yù)熱還有一個副作用就是能把 wasm 實例跑熱后續(xù)的線程調(diào)度也會更穩(wěn)定。5.3 WebGPU 結(jié)果和 CPU 推理不一致心里慌不慌這個現(xiàn)象我遇到的時候也慌了一下。同一張圖CPU 推理出來 Top1 結(jié)果和 WebGPU 推理出來有一點偏差尤其是在 FP16 支持不完整的顯卡上。原因在于 WebGPU 后端為了性能在部分 GPU 算子上會使用半精度浮點計算而 CPU 后端通常用全精度 float32。MobileNet 這類模型對浮點精度誤差有一定容忍度最終特征向量的歸一化數(shù)值不會差很多但如果你在兩個后端上對比精確值會有細微差異。處理方式要分場景。如果是檢索排序微小的特征數(shù)值波動一般不影響 TopK 結(jié)果不用過度擔心。如果你在做嚴格的精度對比實驗可以在創(chuàng)建 session 時把 WebGPU 后端的浮點精度設(shè)置調(diào)整一下。具體參數(shù)在不同版本里命名不一樣建議查 onnxruntime-web 官方文檔。我在項目里用的策略是用搜索結(jié)果的相似度閾值來判斷是否可信低于閾值的查詢直接提示用戶“可能沒有相似圖片”而不是讓用戶盯著一個不靠譜的 Top1 猜來猜去。5.4 跨域與 WASM 加載報錯最容易讓人連夜排查onnxruntime-web 是 WASM 程序要跑多線程還需要 SharedArrayBuffer。瀏覽器對 SharedArrayBuffer 有嚴格限制要求頁面處于跨域隔離狀態(tài)也就是必須帶上Cross-Origin-Opener-Policy: same-origin和Cross-Origin-Embedder-Policy: require-corp這兩個響應(yīng)頭。如果你的頁面部署在普通靜態(tài)服務(wù)器上沒配置這兩個頭控制臺會報一堆跟 WASM 相關(guān)的錯誤最常見的是 “SharedArrayBuffer is not defined” 或者 wasm 加載失敗。很多人會誤以為是模型文件路徑寫錯了其實問題出在響應(yīng)頭。解決方法是讓靜態(tài)服務(wù)器加上響應(yīng)頭同時給模型文件、wasm 文件等所有跨域資源加上Cross-Origin-Resource-Policy: cross-origin。如果你用的是 Nginx配置大概是這樣add_header Cross-Origin-Opener-Policy same-origin always; add_header Cross-Origin-Embedder-Policy require-corp always; add_header Cross-Origin-Resource-Policy cross-origin always;如果模型文件和頁面在同域也要確保模型文件響應(yīng)里帶上 CORP 頭。這個坑很容易出現(xiàn)在“本地開發(fā)好好的一部署到線上就崩”的場景里排查順序應(yīng)該是先看控制臺報錯再檢查響應(yīng)頭最后才懷疑模型文件。5.5 內(nèi)存問題session 重復(fù)創(chuàng)建和 canvas 泄漏瀏覽器端跑推理內(nèi)存泄漏通常不是 GPU 顯存問題而是 JavaScript 堆內(nèi)存被悄悄吃掉。最容易踩的一個坑是每查詢一次就調(diào)用一次InferenceSession.create。創(chuàng)建 session 要加載模型、實例化 wasm、初始化 WebGPU 管線這些對象在多次創(chuàng)建后如果沒有被正確釋放堆內(nèi)存會持續(xù)上漲。我一開始為了“保險”每個查詢都新建 session結(jié)果瀏覽器內(nèi)存從 100 多兆一路漲到 800 多兆頁面開始卡頓。后來改成全局單例 session問題直接消失。另一個坑是 canvas 對象。每次預(yù)處理都創(chuàng)建新的 canvas 并保留在局部變量里沒問題但要避免把 canvas 存在全局數(shù)組或者事件回調(diào)里。用ImageBitmap讀取圖片后要記得調(diào)用close()釋放底層資源尤其是在批量入庫幾百張圖時不釋放的話內(nèi)存漲得飛快。6. 工程化優(yōu)化從能跑到跑得穩(wěn)6.1 會話復(fù)用、預(yù)熱與內(nèi)存回收的落地工程化第一步就是把基礎(chǔ)的資源管理做好。session 全局唯一啟動時創(chuàng)建進程級復(fù)用。預(yù)熱跑一次假推理。索引查詢時給Float32Array用完后及時置空引用讓 GC 可以回收。這里給一個我在項目里用的初始化流程順序很重要checkWebGPU()檢查環(huán)境。創(chuàng)建 session指定executionProviders。執(zhí)行預(yù)熱推理捕獲異常如果 WebGPU 失敗則回退 CPU。從 IndexedDB 加載圖庫特征構(gòu)建內(nèi)存索引。注冊文件上傳、拖拽等事件。初始化是異步過程要給用戶一個 loading 狀態(tài)。瀏覽器端 AI 應(yīng)用最常見的體驗問題就是“沒有任何反饋地白屏幾秒鐘”用戶以為頁面壞了。我習慣在首屏渲染一個進度條至少告訴用戶模型正在加載。6.2 別讓主線程卡死Web Worker 與多請求并行推理和向量檢索都是 CPU/GPU 密集任務(wù)放在主線程會阻塞頁面交互。工程化方案是 Web Worker。在 Worker 里創(chuàng)建 onnxruntime session主線程通過postMessage傳ImageBitmap給 WorkerWorker 完成推理后把特征向量傳回來。ImageBitmap在結(jié)構(gòu)化克隆時是零拷貝的性能影響很小而且 Worker 里跑推理不會卡住頁面。多請求并行的實現(xiàn)要特別注意WebGPU 并不適合無限并發(fā)。GPU 的計算隊列本身是串行調(diào)度的同時發(fā)起幾十個session.run只會增加線程切換開銷并不會加快總吞吐。我的做法是在 Worker 內(nèi)維護一個任務(wù)隊列同一時間最多執(zhí)行一個推理任務(wù)。如果收到多個查詢請求按順序排隊處理。const taskQueue []; let running false; self.onmessage async (event) { taskQueue.push(event.data); if (!running) await processQueue(); }; async function processQueue() { running true; while (taskQueue.length 0) { const task taskQueue.shift(); const imageBitmap task.imageBitmap; const vec await extractFeature(imageBitmap); self.postMessage({ id: task.id, vec: Array.from(vec) }); } running false; }注意postMessage傳Float32Array也是結(jié)構(gòu)化克隆如果向量很大可以改用transferable轉(zhuǎn)移緩沖區(qū)所有權(quán)避免拷貝開銷。但轉(zhuǎn)移之后原數(shù)組就不能用了要確保 Worker 側(cè)不再持有引用。6.3 大圖庫檢索暴力枚舉不是終點純暴力枚舉在幾千張圖的時候表現(xiàn)不錯但圖庫到了十萬張數(shù)量級每查詢一次就要計算十萬次 1024 維內(nèi)積主線程就撐不住了。優(yōu)化思路是聚類剪枝。離線入庫時用 K-Means 把全部特征聚成 K 個類存下每個類的中心向量。在線查詢時先用查詢向量和 K 個中心計算相似度選出最相近的幾個類再只對這幾個類里的特征做精確比對。K 的取值我一般取 64 到 256取決于圖庫總量和目標延遲。這個方案在純前端實現(xiàn)起來不復(fù)雜效果卻非常明顯十萬張圖可以剪到只比對幾千張查詢時間從幾百毫秒降到幾十毫秒。如果向量維度可以降也可以降維。MobileNet 的瓶頸層特征本身有很強的信息冗余我在另一個項目里用 PCA 把 1024 維降到 256 維檢索 Top10 的召回率只下降了 2% 左右但內(nèi)存和計算量直接縮小到原來的四分之一。這個方向?qū)η岸藞鼍昂軐嵱谩?.4 更多落地場景相冊去重、商品同款、素材管理這套方案能做的事不止“以圖搜圖”一個場景。做相冊應(yīng)用時可以批量抽取相冊圖片特征然后做去重和相似聚類用戶不用手動整理幾千張照片。做電商場景時用戶拍一張商品圖瀏覽器直接在本地的商品庫里找同款不上傳圖片購物流程更順暢。做素材管理系統(tǒng)時設(shè)計師上傳一張參考圖本地檢索相似素材響應(yīng)速度遠超云端方案。如果你往里再疊加一個人臉檢測模型做區(qū)域裁剪檢索精度還能進一步提升。我的經(jīng)驗是MobileNet 提取的是全局特征對整圖的配色、構(gòu)圖、物體布局比較敏感但沒法做到精細的局部匹配。遇到這類需求可以在 MobileNet 之前接一個檢測模型或者換用多尺度特征融合的方案別指望一個全局特征解決所有問題。這個工程的延展性其實很好。模型文件可以換后端執(zhí)行器可以切索引策略可以升級但整體的“瀏覽器端特征抽取 本地向量檢索”架構(gòu)是穩(wěn)定可復(fù)用的。只要把extractFeature這個函數(shù)的輸入輸出協(xié)議定好后面換成任何其他特征提取模型業(yè)務(wù)代碼都不用改。這也是我在實際項目中比較受益的一個設(shè)計決策。