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