網(wǎng)絡(luò)訓(xùn)練：從矩陣乘法到反向傳播的完整實(shí)現(xiàn))
1. 從一次“離譜”的需求說起為什么要在瀏覽器里訓(xùn)練神經(jīng)網(wǎng)絡(luò)前陣子有個(gè)做在線教育工具的朋友找我說他們想做一個(gè)“手寫數(shù)字識別”的互動(dòng)教學(xué)頁面學(xué)生用鼠標(biāo)在畫布上寫個(gè)數(shù)字頁面立刻告訴他是幾。聽起來簡單但他們的約束條件很苛刻不能有后端推理服務(wù)不能依賴云端API所有計(jì)算必須在用戶打開網(wǎng)頁的那一瞬間就準(zhǔn)備好而且還得支持“現(xiàn)場教學(xué)微調(diào)”——也就是學(xué)生可以自己補(bǔ)充幾個(gè)樣本讓模型當(dāng)場學(xué)一下識別效果立刻變好。這個(gè)需求翻譯成技術(shù)語言就是在瀏覽器端完成神經(jīng)網(wǎng)絡(luò)的訓(xùn)練和推理。注意是訓(xùn)練不只是推理。推理在瀏覽器里跑早就不新鮮了用ONNX Runtime Web或者TensorFlow.js加載一個(gè)預(yù)訓(xùn)練模型前向傳播一下這事很多年前就能做。但訓(xùn)練意味著反向傳播、梯度計(jì)算、參數(shù)更新計(jì)算量比推理大得多而且要在用戶設(shè)備上跑不能把瀏覽器卡死。我當(dāng)時(shí)的第一個(gè)反應(yīng)是用CPU跑JavaScript別鬧了。一個(gè)兩層全連接網(wǎng)絡(luò)在MNIST上訓(xùn)練純JS循環(huán)跑一輪epoch可能要幾十秒用戶體驗(yàn)直接崩掉。第二個(gè)反應(yīng)是WebAssembly能快一些但WASM本身不直接支持SIMD并行矩陣運(yùn)算除非你自己手寫SIMD指令開發(fā)成本極高。第三個(gè)反應(yīng)也是最終落地的方案WebGL。WebGL的本質(zhì)是什么很多人第一反應(yīng)是“畫3D圖形的”。沒錯(cuò)但更底層地看WebGL是一套在GPU上執(zhí)行并行計(jì)算的接口。它把數(shù)據(jù)組織成紋理把計(jì)算邏輯寫成著色器程序然后讓GPU的幾千個(gè)核心同時(shí)跑。神經(jīng)網(wǎng)絡(luò)的訓(xùn)練核心就是大規(guī)模矩陣乘法和逐元素運(yùn)算這恰好是GPU最擅長的事情。所以用WebGL來加速瀏覽器端神經(jīng)網(wǎng)絡(luò)訓(xùn)練不是“能不能”的問題而是“怎么把神經(jīng)網(wǎng)絡(luò)的計(jì)算圖映射到WebGL的渲染管線上”的問題。這篇文章我就把這條路徑完整拆開講一遍。從GPU為什么適合做這件事到怎么把矩陣乘法寫成著色器到反向傳播怎么在WebGL里落地再到實(shí)際跑起來會(huì)踩哪些坑。如果你正在考慮把模型訓(xùn)練搬到瀏覽器端或者單純好奇“網(wǎng)頁里怎么跑神經(jīng)網(wǎng)絡(luò)”這篇內(nèi)容應(yīng)該能給你一套可復(fù)現(xiàn)的思路。2. GPU的并行邏輯為什么矩陣乘法在WebGL里能快幾十倍2.1 從CPU的“串行思維”切換到GPU的“并行思維”先建立一個(gè)直覺。假設(shè)你要計(jì)算兩個(gè)矩陣相乘A是1000×1000B也是1000×1000結(jié)果C是1000×1000。用CPU做你會(huì)寫三層循環(huán)外層遍歷C的行中層遍歷C的列內(nèi)層做點(diǎn)積?？偣惨?0億次乘加運(yùn)算?，F(xiàn)代CPU大概有8到16個(gè)核心每個(gè)核心同一時(shí)刻只能處理一個(gè)乘加不考慮指令級并行所以理論上要跑幾億個(gè)時(shí)鐘周期。GPU不一樣。一塊普通的集成顯卡也有幾百個(gè)著色器核心獨(dú)立顯卡更是幾千個(gè)。這些核心可以同時(shí)執(zhí)行相同的指令只是處理不同的數(shù)據(jù)。矩陣乘法里C的每一個(gè)元素都是獨(dú)立的互不依賴。這意味著你可以把C的100萬個(gè)元素分配給GPU的幾千個(gè)核心每個(gè)核心算幾百個(gè)元素并行度直接拉滿。但這里有個(gè)關(guān)鍵限制GPU的并行是SIMD式的同一批核心在同一時(shí)刻必須執(zhí)行相同的指令。你不能讓核心A去做乘法核心B去做加法核心C去讀內(nèi)存。所以寫WebGL著色器時(shí)你的思維要變成“我要對每一個(gè)像素或每一個(gè)數(shù)據(jù)元素執(zhí)行同一套操作只是輸入數(shù)據(jù)不同?！?.2 WebGL的渲染管線如何變成計(jì)算管線WebGL的標(biāo)準(zhǔn)用法是你提供頂點(diǎn)數(shù)據(jù)頂點(diǎn)著色器處理每個(gè)頂點(diǎn)然后光柵化片段著色器處理每個(gè)像素最后寫入幀緩沖。這套流程里片段著色器是最適合做通用計(jì)算的階段因?yàn)樗鼘γ總€(gè)像素獨(dú)立執(zhí)行天然并行。要把矩陣乘法映射進(jìn)去做法是這樣的把矩陣A和矩陣B分別編碼成紋理。紋理的每個(gè)像素存儲(chǔ)一個(gè)浮點(diǎn)數(shù)或者用RGBA四個(gè)通道存四個(gè)數(shù)。創(chuàng)建一個(gè)和輸出矩陣C同樣大小的幀緩沖作為渲染目標(biāo)。寫一個(gè)片段著色器對于輸出紋理上的每一個(gè)像素位置(i, j)計(jì)算A的第i行和B的第j列的點(diǎn)積。用一次繪制調(diào)用drawArrays或drawElements觸發(fā)整個(gè)計(jì)算。這里有個(gè)細(xì)節(jié)片段著色器不能直接“讀取”另一個(gè)紋理的任意位置嗎可以用texture2D函數(shù)傳入歸一化的紋理坐標(biāo)。但要注意紋理坐標(biāo)是0到1的浮點(diǎn)數(shù)你需要把整數(shù)索引轉(zhuǎn)換成歸一化坐標(biāo)。比如要讀A的第k列坐標(biāo)就是(k 0.5) / width加0.5是為了采樣到像素中心。一次繪制調(diào)用就能算出整個(gè)C矩陣。1000×1000的矩陣乘法在GPU上就是渲染一個(gè)1000×1000的矩形每個(gè)像素做1000次乘加。GPU的幾千個(gè)核心同時(shí)跑實(shí)際耗時(shí)可能在幾毫秒到幾十毫秒之間比CPU快兩個(gè)數(shù)量級。2.3 精度問題WebGL的浮點(diǎn)數(shù)夠用嗎這是很多人會(huì)忽略的坑。WebGL 1.0的片段著色器默認(rèn)使用mediump精度也就是16位浮點(diǎn)數(shù)范圍大約±65504精度只有10位尾數(shù)。用來做神經(jīng)網(wǎng)絡(luò)訓(xùn)練梯度值很容易下溢到0或者上溢到Infinity。訓(xùn)練直接崩掉。解決方案有兩個(gè)第一在著色器里顯式聲明highp精度。precision highp float;這樣可以用到32位浮點(diǎn)數(shù)具體取決于GPU實(shí)現(xiàn)但大多數(shù)現(xiàn)代設(shè)備支持。highp的范圍和精度足夠覆蓋神經(jīng)網(wǎng)絡(luò)訓(xùn)練中的數(shù)值。第二如果設(shè)備不支持highp一些老舊的移動(dòng)GPU可以考慮用“分塊計(jì)算”或者“縮放因子”來規(guī)避。比如把梯度乘以一個(gè)大的常數(shù)計(jì)算完再除回來。但這會(huì)增加復(fù)雜度實(shí)際項(xiàng)目中我建議直接檢測設(shè)備能力不支持highp就降級到CPU路徑。還有一個(gè)隱藏問題WebGL的紋理格式。默認(rèn)的RGBA紋理每個(gè)通道8位只能存0到255的整數(shù)。要存浮點(diǎn)數(shù)必須使用OES_texture_float擴(kuò)展把紋理內(nèi)部格式設(shè)為gl.FLOAT。這個(gè)擴(kuò)展在現(xiàn)代瀏覽器上支持率很高但還是要做特性檢測。3. 把神經(jīng)網(wǎng)絡(luò)拆成WebGL能吃的“圖形任務(wù)”3.1 全連接層的前向傳播一次矩陣乘法加一次激活假設(shè)你有一個(gè)簡單的全連接層輸入向量x維度n權(quán)重矩陣W維度m×n偏置b維度m輸出y activation(Wx b)。在WebGL里這一步可以拆成兩個(gè)渲染過程第一個(gè)過程算z Wx b。把W編碼成紋理x編碼成另一個(gè)紋理或者直接作為uniform數(shù)組傳入如果n不大的話。輸出z是一個(gè)m維向量可以渲染成一個(gè)m×1的紋理。片段著色器里對于輸出紋理的第i個(gè)像素計(jì)算W的第i行和x的點(diǎn)積再加上b[i]。第二個(gè)過程算y activation(z)。這是一個(gè)逐元素操作直接把z紋理作為輸入渲染到y(tǒng)紋理。著色器里對每個(gè)像素應(yīng)用激活函數(shù)比如ReLU就是max(0, z)Sigmoid就是1/(1exp(-z))。這里有個(gè)優(yōu)化點(diǎn)如果網(wǎng)絡(luò)有多層可以把中間結(jié)果一直保留在紋理里不需要讀回CPU。每一層的輸出紋理直接作為下一層的輸入紋理。這樣整個(gè)前向傳播就是一系列繪制調(diào)用數(shù)據(jù)始終在GPU顯存里流轉(zhuǎn)避免了PCIe總線的來回拷貝。3.2 反向傳播的鏈?zhǔn)椒▌t在著色器里怎么表達(dá)反向傳播是訓(xùn)練的核心也是最難映射到WebGL的部分。核心公式是對于每一層先算輸出的梯度dL/dz然后算權(quán)重的梯度dL/dW dL/dz * x^T輸入的梯度dL/dx W^T * dL/dz。在WebGL里dL/dz的計(jì)算通常涉及逐元素乘法比如ReLU的導(dǎo)數(shù)就是z0?1:0和dL/dy逐元素相乘。這個(gè)可以寫成一個(gè)片段著色器輸入是dL/dy紋理和z紋理輸出是dL/dz紋理。dL/dW的計(jì)算是一個(gè)外積dL/dz是m維x是n維外積得到m×n的矩陣。這個(gè)可以映射成一次矩陣乘法把dL/dz當(dāng)作列向量x當(dāng)作行向量。輸出紋理的大小是m×n每個(gè)像素(i,j)的值是dL/dz[i] * x[j]。這個(gè)計(jì)算很簡單但要注意紋理的布局——WebGL紋理是二維的你需要決定怎么把一維向量和二維矩陣對應(yīng)起來。dL/dx的計(jì)算是W^T乘以dL/dz這是一個(gè)矩陣向量乘法。W^T是n×mdL/dz是m維結(jié)果是n維。這個(gè)可以復(fù)用前面矩陣乘法的著色器只是把W的紋理轉(zhuǎn)置一下或者在著色器里交換索引。實(shí)際實(shí)現(xiàn)時(shí)我建議把每一層的反向傳播拆成三個(gè)獨(dú)立的渲染過程算dL/dz、算dL/dW、算dL/dx。每個(gè)過程一個(gè)著色器程序輸入輸出都是紋理。這樣邏輯清晰調(diào)試也方便。3.3 參數(shù)更新在GPU上做還是回CPU做參數(shù)更新是W W - learning_rate * dL/dW。這個(gè)操作本身很簡單逐元素減法。但問題在于更新后的W需要作為下一輪前向傳播的輸入。如果W一直留在GPU紋理里那更新也可以在GPU上做寫一個(gè)片段著色器輸入W紋理和dL/dW紋理輸出新的W紋理。但這里有個(gè)權(quán)衡。如果所有參數(shù)都在GPU上你就沒法方便地查看訓(xùn)練過程中的權(quán)重變化也沒法做復(fù)雜的優(yōu)化器比如Adam需要維護(hù)一階矩和二階矩計(jì)算量不大但邏輯稍復(fù)雜。我的做法是前向和反向在GPU上跑參數(shù)更新回CPU做。每一輪迭代結(jié)束后把dL/dW從GPU讀回CPU用readPixels在JavaScript里更新權(quán)重再把新權(quán)重上傳到GPU紋理。這樣雖然多了兩次數(shù)據(jù)傳輸?shù)`活性大大提高而且對于小網(wǎng)絡(luò)來說傳輸開銷遠(yuǎn)小于計(jì)算開銷。讀回?cái)?shù)據(jù)時(shí)要注意readPixels是同步操作會(huì)阻塞渲染管線。如果每輪都讀可能會(huì)成為瓶頸。優(yōu)化方法是批量讀——積累幾輪的梯度再一起讀回或者用WebGL 2.0的PBO像素緩沖對象做異步讀回。但WebGL 2.0的兼容性還不如1.0實(shí)際項(xiàng)目中要權(quán)衡。4. 手寫一個(gè)WebGL神經(jīng)網(wǎng)絡(luò)訓(xùn)練器的完整路徑4.1 環(huán)境搭建與基礎(chǔ)工具函數(shù)首先你需要一個(gè)WebGL上下文。創(chuàng)建一個(gè)canvas獲取gl對象。然后檢查必要的擴(kuò)展const canvas document.createElement(canvas); const gl canvas.getContext(webgl); const floatExt gl.getExtension(OES_texture_float); if (!floatExt) { throw new Error(設(shè)備不支持浮點(diǎn)紋理無法進(jìn)行GPU訓(xùn)練); }接下來是幾個(gè)基礎(chǔ)工具函數(shù)這些函數(shù)會(huì)反復(fù)用到創(chuàng)建浮點(diǎn)紋理輸入是一個(gè)Float32Array輸出是一個(gè)WebGL紋理對象。注意設(shè)置紋理參數(shù)為NEAREST采樣避免插值導(dǎo)致數(shù)值錯(cuò)誤。function createFloatTexture(gl, data, width, height) { const texture gl.createTexture(); gl.bindTexture(gl.TEXTURE_2D, texture); gl.texImage2D(gl.TEXTURE_2D, 0, gl.RGBA, width, height, 0, gl.RGBA, gl.FLOAT, data); gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MIN_FILTER, gl.NEAREST); gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MAG_FILTER, gl.NEAREST); gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_S, gl.CLAMP_TO_EDGE); gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_T, gl.CLAMP_TO_EDGE); return texture; }創(chuàng)建幀緩沖幀緩沖是渲染目標(biāo)你把片段著色器的輸出寫進(jìn)去。每個(gè)幀緩沖綁定一個(gè)紋理作為顏色附件。function createFramebuffer(gl, texture) { const fb gl.createFramebuffer(); gl.bindFramebuffer(gl.FRAMEBUFFER, fb); gl.framebufferTexture2D(gl.FRAMEBUFFER, gl.COLOR_ATTACHMENT0, gl.TEXTURE_2D, texture, 0); return fb; }編譯著色器程序這是標(biāo)準(zhǔn)流程編譯頂點(diǎn)著色器和片段著色器鏈接成程序。頂點(diǎn)著色器通常很簡單就是傳遞一個(gè)覆蓋整個(gè)畫布的矩形。function createProgram(gl, vertexSource, fragmentSource) { const vs gl.createShader(gl.VERTEX_SHADER); gl.shaderSource(vs, vertexSource); gl.compileShader(vs); const fs gl.createShader(gl.FRAGMENT_SHADER); gl.shaderSource(fs, fragmentSource); gl.compileShader(fs); const program gl.createProgram(); gl.attachShader(program, vs); gl.attachShader(program, fs); gl.linkProgram(program); return program; }這些工具函數(shù)看起來簡單但每個(gè)都有坑。比如紋理的寬高必須是2的冪嗎在WebGL 1.0里如果使用NEAREST采樣和CLAMP_TO_EDGE可以不是2的冪。但有些驅(qū)動(dòng)會(huì)強(qiáng)制要求所以最安全的做法是padding到2的冪。我一般會(huì)把矩陣維度補(bǔ)齊到最近的2的冪比如1000×1000的矩陣實(shí)際分配1024×1024的紋理多出來的部分填0。4.2 矩陣乘法著色器的逐行拆解矩陣乘法是整個(gè)訓(xùn)練流程里最核心的計(jì)算。假設(shè)我們要算C A × BA是M×KB是K×NC是M×N。頂點(diǎn)著色器很簡單就是畫一個(gè)覆蓋整個(gè)輸出紋理的矩形attribute vec2 a_position; void main() { gl_Position vec4(a_position, 0.0, 1.0); }片段著色器是重點(diǎn)precision highp float; uniform sampler2D u_A; uniform sampler2D u_B; uniform vec2 u_A_size; // (M, K) uniform vec2 u_B_size; // (K, N) uniform vec2 u_C_size; // (M, N) void main() { vec2 c_pos gl_FragCoord.xy; // 當(dāng)前像素在C中的位置 float sum 0.0; for (int k 0; k int(u_A_size.y); k) { // 讀A的第c_pos.x行第k列 vec2 a_uv vec2((c_pos.x 0.5) / u_A_size.x, (float(k) 0.5) / u_A_size.y); float a_val texture2D(u_A, a_uv).r; // 讀B的第k行第c_pos.y列 vec2 b_uv vec2((float(k) 0.5) / u_B_size.x, (c_pos.y 0.5) / u_B_size.y); float b_val texture2D(u_B, b_uv).r; sum a_val * b_val; } gl_FragColor vec4(sum, 0.0, 0.0, 1.0); }這段代碼有幾個(gè)關(guān)鍵點(diǎn)第一gl_FragCoord.xy給出的是像素中心坐標(biāo)比如第一個(gè)像素是(0.5, 0.5)。所以計(jì)算紋理坐標(biāo)時(shí)要加0.5再除以尺寸。第二循環(huán)次數(shù)必須是常量或者可以用uniform控制。WebGL 1.0的GLSL不允許非常量循環(huán)邊界但可以用一個(gè)足夠大的常量循環(huán)然后在循環(huán)體內(nèi)用if判斷是否超出范圍。不過現(xiàn)代瀏覽器對動(dòng)態(tài)循環(huán)的支持已經(jīng)不錯(cuò)了實(shí)際測試中大多數(shù)設(shè)備都能跑。第三每次循環(huán)要讀兩次紋理。紋理讀取是GPU計(jì)算中最慢的操作之一因?yàn)橐L問顯存。優(yōu)化方法是把A和B的紋理綁定到不同的紋理單元并且盡量讓循環(huán)內(nèi)的訪問模式有局部性。但矩陣乘法的訪問模式天然是跳躍的所以這個(gè)優(yōu)化空間有限。一個(gè)實(shí)用的優(yōu)化技巧分塊計(jì)算。把大矩陣切成小塊比如32×32的塊每個(gè)塊的計(jì)算可以更好地利用GPU的緩存。但這會(huì)增加著色器的復(fù)雜度實(shí)際項(xiàng)目中如果矩陣不大比如全連接層通常幾百維直接算就夠了。4.3 激活函數(shù)與損失函數(shù)的GPU實(shí)現(xiàn)激活函數(shù)是逐元素操作實(shí)現(xiàn)起來簡單。以ReLU為例precision highp float; uniform sampler2D u_input; void main() { vec2 uv gl_FragCoord.xy / vec2(textureSize(u_input, 0)); float val texture2D(u_input, uv).r; gl_FragColor vec4(max(0.0, val), 0.0, 0.0, 1.0); }Sigmoid稍微復(fù)雜一點(diǎn)因?yàn)橐鉫xpfloat sigmoid(float x) { return 1.0 / (1.0 exp(-x)); }但要注意exp函數(shù)在highp下可能會(huì)有精度問題特別是當(dāng)x很大或很小時(shí)。一個(gè)穩(wěn)定的實(shí)現(xiàn)是float stable_sigmoid(float x) { if (x 0.0) { return 1.0 / (1.0 exp(-x)); } else { float e exp(x); return e / (1.0 e); } }損失函數(shù)通常用交叉熵。對于分類問題交叉熵的梯度是softmax輸出減去真實(shí)標(biāo)簽。這個(gè)計(jì)算可以放在CPU做因?yàn)橹簧婕拜敵鰧佑?jì)算量小。但如果想全GPU化也可以寫一個(gè)著色器輸入是預(yù)測概率和真實(shí)標(biāo)簽輸出是每個(gè)樣本的損失值。4.4 反向傳播的著色器實(shí)現(xiàn)細(xì)節(jié)反向傳播里最復(fù)雜的是dL/dW的計(jì)算。假設(shè)當(dāng)前層的輸入是xn維輸出梯度是dL/dzm維那么dL/dW是一個(gè)m×n的矩陣每個(gè)元素是dL/dz[i] * x[j]。這個(gè)計(jì)算可以寫成一個(gè)著色器輸出紋理大小是m×nprecision highp float; uniform sampler2D u_dz; // m維向量存成m×1紋理 uniform sampler2D u_x; // n維向量存成n×1紋理 uniform vec2 u_dz_size; uniform vec2 u_x_size; void main() { vec2 pos gl_FragCoord.xy; // pos.x是行索引ipos.y是列索引j vec2 dz_uv vec2(0.5 / u_dz_size.x, (pos.x 0.5) / u_dz_size.y); vec2 x_uv vec2(0.5 / u_x_size.x, (pos.y 0.5) / u_x_size.y); float dz_val texture2D(u_dz, dz_uv).r; float x_val texture2D(u_x, x_uv).r; gl_FragColor vec4(dz_val * x_val, 0.0, 0.0, 1.0); }注意這里紋理坐標(biāo)的寫法。因?yàn)閐z和x都是一維向量我習(xí)慣把它們存成寬度為1、高度為維度的紋理。這樣讀取時(shí)x坐標(biāo)固定為0.5/width即像素中心y坐標(biāo)根據(jù)索引變化。dL/dx的計(jì)算是W^T × dL/dz。W是m×nW^T是n×mdL/dz是m維結(jié)果是n維。這個(gè)可以復(fù)用矩陣乘法的著色器只需要把W的紋理坐標(biāo)交換一下?；蛘邔懸粋€(gè)專門的著色器precision highp float; uniform sampler2D u_W; // m×n uniform sampler2D u_dz; // m維 uniform vec2 u_W_size; uniform vec2 u_dz_size; void main() { float i gl_FragCoord.y; // 輸出向量的索引 float sum 0.0; for (int j 0; j int(u_W_size.x); j) { vec2 w_uv vec2((float(j) 0.5) / u_W_size.x, (i 0.5) / u_W_size.y); float w_val texture2D(u_W, w_uv).r; vec2 dz_uv vec2(0.5 / u_dz_size.x, (float(j) 0.5) / u_dz_size.y); float dz_val texture2D(u_dz, dz_uv).r; sum w_val * dz_val; } gl_FragColor vec4(sum, 0.0, 0.0, 1.0); }這里W的讀取是W[j][i]因?yàn)槲覀円鉝^T的第i行也就是W的第i列。所以w_uv的x坐標(biāo)是j行索引y坐標(biāo)是i列索引。這個(gè)索引交換很容易搞混寫的時(shí)候要特別小心。5. 實(shí)測中那些文檔不會(huì)告訴你的坑5.1 紋理尺寸限制與矩陣分塊策略WebGL有一個(gè)硬性限制MAX_TEXTURE_SIZE。這個(gè)值在不同設(shè)備上不一樣低端手機(jī)可能只有2048高端顯卡能到16384。如果你的權(quán)重矩陣超過這個(gè)尺寸就必須分塊。分塊矩陣乘法的思路是把A切成若干行塊B切成若干列塊然后逐塊計(jì)算。每個(gè)塊的計(jì)算結(jié)果寫到一個(gè)臨時(shí)紋理最后拼起來。但這樣會(huì)增加繪制調(diào)用次數(shù)和紋理切換開銷。我的經(jīng)驗(yàn)是對于全連接層權(quán)重矩陣通常不會(huì)超過4096×4096那已經(jīng)是1600萬參數(shù)了所以大多數(shù)設(shè)備都能直接放下。但如果做卷積層特征圖可能很大就需要分塊。實(shí)際項(xiàng)目中我建議先檢測MAX_TEXTURE_SIZE如果矩陣維度超過限制就自動(dòng)降級到CPU計(jì)算或者用分塊策略。還有一個(gè)隱藏限制MAX_RENDERBUFFER_SIZE和MAX_VIEWPORT_DIMS。這些通常和紋理尺寸限制一致但有些設(shè)備上幀緩沖的尺寸限制更嚴(yán)格。創(chuàng)建幀緩沖后一定要檢查checkFramebufferStatus確保返回FRAMEBUFFER_COMPLETE。5.2 浮點(diǎn)紋理的讀寫精度陷阱前面提到要用OES_texture_float擴(kuò)展。但即使擴(kuò)展存在不同設(shè)備對浮點(diǎn)紋理的支持程度也不一樣。有些設(shè)備只能渲染到浮點(diǎn)紋理但不能讀取有些設(shè)備讀取時(shí)會(huì)有精度損失。一個(gè)實(shí)用的檢測方法是創(chuàng)建一個(gè)浮點(diǎn)紋理寫入一個(gè)已知值比如0.123456789然后讀回來看精度是否滿足要求。如果誤差超過1e-6就說明這個(gè)設(shè)備的浮點(diǎn)紋理不可靠需要降級。另外readPixels讀回浮點(diǎn)數(shù)據(jù)時(shí)返回的是Float32Array。但有些設(shè)備會(huì)返回Uint8Array即使你請求的是浮點(diǎn)格式。所以讀回后要檢查數(shù)據(jù)類型必要時(shí)做轉(zhuǎn)換。還有一個(gè)坑WebGL的texImage2D在傳入Float32Array時(shí)如果數(shù)據(jù)長度和width×height×4不匹配會(huì)靜默失敗或者產(chǎn)生未定義行為。一定要確保數(shù)據(jù)長度正確。我習(xí)慣在創(chuàng)建紋理前先算好width * height * 4然后檢查數(shù)組長度。5.3 訓(xùn)練不收斂的排查鏈路GPU訓(xùn)練最讓人頭疼的問題是訓(xùn)練不收斂但你看不到中間過程。CPU訓(xùn)練時(shí)你可以打印每一層的梯度、權(quán)重但GPU上數(shù)據(jù)都在紋理里讀回來很麻煩。我總結(jié)了一套排查鏈路第一步檢查前向傳播。用一個(gè)已知的輸入手動(dòng)算一遍CPU版本的前向結(jié)果然后和GPU版本對比。如果前向就不對那反向肯定不對。對比時(shí)要注意精度GPU的浮點(diǎn)運(yùn)算和CPU可能有微小差異但應(yīng)該在1e-5以內(nèi)。第二步檢查梯度。用數(shù)值梯度法驗(yàn)證對于某個(gè)權(quán)重w計(jì)算Loss(wε)和Loss(w-ε)然后(Loss(wε)-Loss(w-ε))/(2ε)就是數(shù)值梯度。和GPU算出的梯度對比如果差異大說明反向傳播有問題。第三步檢查參數(shù)更新。如果梯度對了但Loss不下降可能是學(xué)習(xí)率太大或太小。GPU訓(xùn)練時(shí)梯度值可能因?yàn)榫葐栴}被截?cái)鄬?dǎo)致實(shí)際學(xué)習(xí)率和設(shè)定值不符?？梢韵劝褜W(xué)習(xí)率調(diào)小10倍試試。第四步檢查數(shù)據(jù)。有時(shí)候問題不在GPU而在輸入數(shù)據(jù)。比如標(biāo)簽編碼錯(cuò)了或者輸入沒有歸一化。GPU訓(xùn)練對輸入尺度很敏感因?yàn)楦↑c(diǎn)數(shù)的動(dòng)態(tài)范圍有限。5.4 性能調(diào)優(yōu)從每秒幾輪到每秒幾百輪最初的實(shí)現(xiàn)可能每秒只能跑幾輪迭代慢得讓人懷疑人生。優(yōu)化方向有幾個(gè)減少繪制調(diào)用。每次繪制都有固定開銷包括狀態(tài)切換、著色器綁定、紋理綁定。如果能把多個(gè)操作合并到一個(gè)著色器里就能減少調(diào)用次數(shù)。比如把矩陣乘法和偏置加法合并把激活函數(shù)和損失函數(shù)合并。使用頂點(diǎn)緩沖對象VBO。畫矩形時(shí)不要每次都用gl.vertexAttribPointer重新指定頂點(diǎn)數(shù)據(jù)。創(chuàng)建一個(gè)VBO把矩形的四個(gè)頂點(diǎn)存進(jìn)去每次繪制時(shí)直接綁定。避免不必要的紋理切換。紋理切換是昂貴的操作。如果多個(gè)著色器使用相同的輸入紋理盡量把它們安排在連續(xù)的繪制調(diào)用里。使用WebGL 2.0。WebGL 2.0支持更多紋理格式、多重渲染目標(biāo)MRT、以及更靈活的循環(huán)控制。如果目標(biāo)用戶設(shè)備支持優(yōu)先用WebGL 2.0。但要注意兼容性做好降級方案。批處理。如果訓(xùn)練數(shù)據(jù)是小樣本可以把多個(gè)樣本打包成一個(gè)批次一次前向傳播算完。這樣GPU的利用率更高因?yàn)槊總€(gè)著色器調(diào)用處理的數(shù)據(jù)量更大。實(shí)測下來一個(gè)兩層全連接網(wǎng)絡(luò)784-128-10在MNIST上訓(xùn)練優(yōu)化前每秒約5輪優(yōu)化后能到每秒200輪以上。這個(gè)速度已經(jīng)足夠做交互式教學(xué)了——學(xué)生寫一個(gè)數(shù)字模型在幾百毫秒內(nèi)就能完成一輪微調(diào)。6. 瀏覽器端訓(xùn)練的邊界與適用場景6.1 什么規(guī)模的模型適合放在瀏覽器里訓(xùn)不是所有模型都適合瀏覽器端訓(xùn)練。我的經(jīng)驗(yàn)是參數(shù)量在10萬到100萬之間的模型WebGL訓(xùn)練是可行的。超過這個(gè)規(guī)模紋理尺寸和顯存會(huì)成為瓶頸而且訓(xùn)練時(shí)間會(huì)變得不可接受。具體來說全連接層每層幾百到幾千維層數(shù)不超過5層這樣的網(wǎng)絡(luò)在瀏覽器里訓(xùn)練很流暢。卷積層也可以做但需要把卷積展開成矩陣乘法im2col這會(huì)增加內(nèi)存占用。如果卷積核不大比如3×3通道數(shù)不多比如32或64也是可以跑的。不適合的場景包括大型語言模型、高分辨率圖像生成模型、需要大量迭代才能收斂的模型。這些還是交給專業(yè)訓(xùn)練環(huán)境。6.2 隱私保護(hù)與離線可用性的真實(shí)價(jià)值瀏覽器端訓(xùn)練最大的賣點(diǎn)不是性能而是數(shù)據(jù)不出本地。用戶的輸入數(shù)據(jù)、標(biāo)注數(shù)據(jù)、甚至模型本身都可以完全留在瀏覽器里。這對于隱私敏感的應(yīng)用場景很有價(jià)值比如醫(yī)療健康類的個(gè)人數(shù)據(jù)訓(xùn)練、教育類的學(xué)生行為數(shù)據(jù)建模。另一個(gè)價(jià)值是離線可用。一旦頁面加載完成后續(xù)的訓(xùn)練和推理都不需要網(wǎng)絡(luò)連接。這對于網(wǎng)絡(luò)不穩(wěn)定的環(huán)境或者需要完全離線部署的場景是很實(shí)用的特性。還有一個(gè)容易被忽略的價(jià)值即時(shí)反饋。用戶補(bǔ)充幾個(gè)樣本模型立刻更新識別效果當(dāng)場變好。這種交互體驗(yàn)是云端訓(xùn)練做不到的因?yàn)樵贫擞?xùn)練有網(wǎng)絡(luò)延遲和排隊(duì)等待。6.3 和TensorFlow.js的GPU后端有什么不同TensorFlow.js也支持WebGL后端而且封裝得很好幾行代碼就能定義模型并訓(xùn)練。那為什么還要自己寫WebGL著色器區(qū)別在于控制粒度。TensorFlow.js的WebGL后端是一個(gè)通用框架它會(huì)把各種操作映射到著色器但中間有很多抽象層。對于標(biāo)準(zhǔn)操作比如全連接、卷積、常見的激活函數(shù)它的性能已經(jīng)很好。但如果你需要自定義操作或者想針對特定硬件做極致優(yōu)化自己寫著色器更靈活。另外TensorFlow.js的WebGL后端在某些設(shè)備上會(huì)有兼容性問題而且它的內(nèi)存管理是自動(dòng)的有時(shí)候會(huì)導(dǎo)致紋理泄漏或者頻繁的垃圾回收。自己管理紋理和幀緩沖雖然代碼量大但行為完全可控。我的建議是如果標(biāo)準(zhǔn)模型能滿足需求直接用TensorFlow.js省時(shí)省力。如果需要自定義訓(xùn)練邏輯或者對性能有極致要求再考慮手寫WebGL。7. 從原型到產(chǎn)品工程化落地的幾個(gè)關(guān)鍵決策7.1 如何設(shè)計(jì)降級方案保證兼容性不是所有瀏覽器都支持WebGL也不是所有支持WebGL的設(shè)備都支持浮點(diǎn)紋理。產(chǎn)品化時(shí)必須設(shè)計(jì)降級方案。我的做法是三級降級第一級WebGL 2.0 浮點(diǎn)紋理。這是最優(yōu)路徑性能最好。第二級WebGL 1.0 浮點(diǎn)紋理擴(kuò)展。性能稍差但兼容性更廣。第三級純CPU JavaScript。用TypedArray做矩陣運(yùn)算雖然慢但保證能用。對于小模型參數(shù)量小于1萬CPU版本也能在可接受的時(shí)間內(nèi)完成訓(xùn)練。降級檢測在頁面加載時(shí)做一次把結(jié)果緩存起來。用戶不需要知道底層用的是哪條路徑只需要看到訓(xùn)練能跑起來。7.2 訓(xùn)練過程的進(jìn)度反饋與中斷處理GPU訓(xùn)練是異步的你發(fā)出繪制調(diào)用后GPU在后臺跑JavaScript繼續(xù)執(zhí)行。如果訓(xùn)練循環(huán)寫成一個(gè)大的for循環(huán)頁面會(huì)卡死用戶看不到任何進(jìn)度。解決方案是用requestAnimationFrame或者setTimeout把訓(xùn)練循環(huán)拆成小批次。每一幀跑幾輪迭代然后更新進(jìn)度條再請求下一幀。這樣頁面保持響應(yīng)用戶能看到訓(xùn)練進(jìn)度。中斷處理也很重要。用戶可能中途關(guān)閉頁面或者切換標(biāo)簽頁。要在beforeunload事件里清理WebGL資源避免內(nèi)存泄漏。另外如果用戶切換標(biāo)簽頁requestAnimationFrame會(huì)暫停訓(xùn)練也會(huì)暫停。這是符合預(yù)期的行為但要在UI上給出提示。7.3 模型持久化把訓(xùn)練好的權(quán)重存到哪里訓(xùn)練完成后權(quán)重需要保存。GPU紋理里的數(shù)據(jù)不能直接序列化必須先讀回CPU轉(zhuǎn)成Float32Array然后存到IndexedDB或者localStorage。IndexedDB適合存大文件幾MB到幾十MBlocalStorage只適合存小數(shù)據(jù)幾KB。對于神經(jīng)網(wǎng)絡(luò)權(quán)重通常用IndexedDB。存儲(chǔ)時(shí)可以把Float32Array轉(zhuǎn)成Blob或者直接存ArrayBuffer。讀取時(shí)反向操作從IndexedDB讀出ArrayBuffer轉(zhuǎn)成Float32Array再上傳到GPU紋理。整個(gè)過程要注意異步處理因?yàn)镮ndexedDB的讀寫都是異步的。還有一個(gè)細(xì)節(jié)權(quán)重的版本管理。如果模型結(jié)構(gòu)變了舊權(quán)重就不能用了。我習(xí)慣在存儲(chǔ)時(shí)帶上一個(gè)版本號或者結(jié)構(gòu)哈希讀取時(shí)校驗(yàn)不匹配就重新初始化。8. 我踩過的三個(gè)印象最深的坑第一個(gè)坑是紋理坐標(biāo)的0.5偏移。剛開始寫著色器時(shí)我直接用了gl_FragCoord.xy / textureSize作為紋理坐標(biāo)結(jié)果采樣到的總是相鄰像素的值導(dǎo)致矩陣乘法結(jié)果完全錯(cuò)誤。排查了很久才發(fā)現(xiàn)gl_FragCoord給出的是像素中心坐標(biāo)而紋理坐標(biāo)需要的是歸一化的像素中心所以要加0.5再除。這個(gè)細(xì)節(jié)在大多數(shù)WebGL教程里都不會(huì)強(qiáng)調(diào)但做通用計(jì)算時(shí)是致命的。第二個(gè)坑是浮點(diǎn)紋理的線性插值。默認(rèn)情況下紋理采樣使用線性插值會(huì)混合相鄰像素的值。對于圖像渲染這是好事但對于矩陣運(yùn)算這是災(zāi)難。必須把TEXTURE_MIN_FILTER和TEXTURE_MAG_FILTER都設(shè)為NEAREST。我一開始忘了設(shè)結(jié)果梯度計(jì)算出來總是有微小的噪聲訓(xùn)練完全不收斂。第三個(gè)坑是幀緩沖的完整性檢查。創(chuàng)建幀緩沖后我沒有檢查checkFramebufferStatus直接就開始渲染。在某些設(shè)備上幀緩沖創(chuàng)建失敗但不會(huì)報(bào)錯(cuò)渲染結(jié)果全是0。后來養(yǎng)成了習(xí)慣每次創(chuàng)建幀緩沖后都檢查狀態(tài)不完整就拋異常。這個(gè)檢查在開發(fā)階段能省下大量調(diào)試時(shí)間。9. 后續(xù)可以繼續(xù)深挖的方向如果你已經(jīng)把基礎(chǔ)的WebGL訓(xùn)練跑通了還有幾個(gè)方向可以繼續(xù)優(yōu)化。一是用WebGL 2.0的變換反饋Transform Feedback。WebGL 2.0允許把頂點(diǎn)著色器的輸出直接寫回緩沖區(qū)不需要經(jīng)過片段著色器。這對于逐元素操作比如激活函數(shù)、參數(shù)更新可能更高效因?yàn)楸苊饬思y理的創(chuàng)建和綁定。二是混合精度訓(xùn)練。用16位浮點(diǎn)數(shù)做前向和反向用32位浮點(diǎn)數(shù)做參數(shù)更新。這樣可以減少顯存占用和帶寬消耗速度可能提升一倍。但要注意數(shù)值穩(wěn)定性可能需要損失縮放loss scaling來防止梯度下溢。三是多GPU協(xié)同。如果用戶設(shè)備有多個(gè)GPU比如集成顯卡加獨(dú)立顯卡理論上可以用多個(gè)WebGL上下文并行計(jì)算。但瀏覽器對多GPU的支持還很有限實(shí)際可行性不高。四是和WebAssembly結(jié)合。把一些不適合GPU的操作比如數(shù)據(jù)預(yù)處理、復(fù)雜的優(yōu)化器邏輯放到WASM里做GPU專注做矩陣運(yùn)算。兩者通過SharedArrayBuffer通信實(shí)現(xiàn)真正的并行。這條路我走了大半年從最初的“這能行嗎”到現(xiàn)在的“基本能用”中間踩了無數(shù)坑但也積累了不少經(jīng)驗(yàn)。瀏覽器端訓(xùn)練不是要取代專業(yè)訓(xùn)練環(huán)境而是在特定場景下提供一種輕量、隱私、即時(shí)的選擇。如果你的項(xiàng)目恰好需要這些特性WebGL方案值得一試。 WebGL加速瀏覽器端神經(jīng)網(wǎng)絡(luò)訓(xùn)練：從矩陣乘法到反向傳播的完整實(shí)現(xiàn)
1. 從一次“離譜”的需求說起為什么要在瀏覽器里訓(xùn)練神經(jīng)網(wǎng)絡(luò)前陣子有個(gè)做在線教育工具的朋友找我說他們想做一個(gè)“手寫數(shù)字識別”的互動(dòng)教學(xué)頁面學(xué)生用鼠標(biāo)在畫布上寫個(gè)數(shù)字頁面立刻告訴他是幾。聽起來簡單但他們的約束條件很苛刻不能有后端推理服務(wù)不能依賴云端API所有計(jì)算必須在用戶打開網(wǎng)頁的那一瞬間就準(zhǔn)備好而且還得支持“現(xiàn)場教學(xué)微調(diào)”——也就是學(xué)生可以自己補(bǔ)充幾個(gè)樣本讓模型當(dāng)場學(xué)一下識別效果立刻變好。這個(gè)需求翻譯成技術(shù)語言就是在瀏覽器端完成神經(jīng)網(wǎng)絡(luò)的訓(xùn)練和推理。注意是訓(xùn)練不只是推理。推理在瀏覽器里跑早就不新鮮了用ONNX Runtime Web或者TensorFlow.js加載一個(gè)預(yù)訓(xùn)練模型前向傳播一下這事很多年前就能做。但訓(xùn)練意味著反向傳播、梯度計(jì)算、參數(shù)更新計(jì)算量比推理大得多而且要在用戶設(shè)備上跑不能把瀏覽器卡死。我當(dāng)時(shí)的第一個(gè)反應(yīng)是用CPU跑JavaScript別鬧了。一個(gè)兩層全連接網(wǎng)絡(luò)在MNIST上訓(xùn)練純JS循環(huán)跑一輪epoch可能要幾十秒用戶體驗(yàn)直接崩掉。第二個(gè)反應(yīng)是WebAssembly能快一些但WASM本身不直接支持SIMD并行矩陣運(yùn)算除非你自己手寫SIMD指令開發(fā)成本極高。第三個(gè)反應(yīng)也是最終落地的方案WebGL。WebGL的本質(zhì)是什么很多人第一反應(yīng)是“畫3D圖形的”。沒錯(cuò)但更底層地看WebGL是一套在GPU上執(zhí)行并行計(jì)算的接口。它把數(shù)據(jù)組織成紋理把計(jì)算邏輯寫成著色器程序然后讓GPU的幾千個(gè)核心同時(shí)跑。神經(jīng)網(wǎng)絡(luò)的訓(xùn)練核心就是大規(guī)模矩陣乘法和逐元素運(yùn)算這恰好是GPU最擅長的事情。所以用WebGL來加速瀏覽器端神經(jīng)網(wǎng)絡(luò)訓(xùn)練不是“能不能”的問題而是“怎么把神經(jīng)網(wǎng)絡(luò)的計(jì)算圖映射到WebGL的渲染管線上”的問題。這篇文章我就把這條路徑完整拆開講一遍。從GPU為什么適合做這件事到怎么把矩陣乘法寫成著色器到反向傳播怎么在WebGL里落地再到實(shí)際跑起來會(huì)踩哪些坑。如果你正在考慮把模型訓(xùn)練搬到瀏覽器端或者單純好奇“網(wǎng)頁里怎么跑神經(jīng)網(wǎng)絡(luò)”這篇內(nèi)容應(yīng)該能給你一套可復(fù)現(xiàn)的思路。2. GPU的并行邏輯為什么矩陣乘法在WebGL里能快幾十倍2.1 從CPU的“串行思維”切換到GPU的“并行思維”先建立一個(gè)直覺。假設(shè)你要計(jì)算兩個(gè)矩陣相乘A是1000×1000B也是1000×1000結(jié)果C是1000×1000。用CPU做你會(huì)寫三層循環(huán)外層遍歷C的行中層遍歷C的列內(nèi)層做點(diǎn)積?？偣惨?0億次乘加運(yùn)算?，F(xiàn)代CPU大概有8到16個(gè)核心每個(gè)核心同一時(shí)刻只能處理一個(gè)乘加不考慮指令級并行所以理論上要跑幾億個(gè)時(shí)鐘周期。GPU不一樣。一塊普通的集成顯卡也有幾百個(gè)著色器核心獨(dú)立顯卡更是幾千個(gè)。這些核心可以同時(shí)執(zhí)行相同的指令只是處理不同的數(shù)據(jù)。矩陣乘法里C的每一個(gè)元素都是獨(dú)立的互不依賴。這意味著你可以把C的100萬個(gè)元素分配給GPU的幾千個(gè)核心每個(gè)核心算幾百個(gè)元素并行度直接拉滿。但這里有個(gè)關(guān)鍵限制GPU的并行是SIMD式的同一批核心在同一時(shí)刻必須執(zhí)行相同的指令。你不能讓核心A去做乘法核心B去做加法核心C去讀內(nèi)存。所以寫WebGL著色器時(shí)你的思維要變成“我要對每一個(gè)像素或每一個(gè)數(shù)據(jù)元素執(zhí)行同一套操作只是輸入數(shù)據(jù)不同?！?.2 WebGL的渲染管線如何變成計(jì)算管線WebGL的標(biāo)準(zhǔn)用法是你提供頂點(diǎn)數(shù)據(jù)頂點(diǎn)著色器處理每個(gè)頂點(diǎn)然后光柵化片段著色器處理每個(gè)像素最后寫入幀緩沖。這套流程里片段著色器是最適合做通用計(jì)算的階段因?yàn)樗鼘γ總€(gè)像素獨(dú)立執(zhí)行天然并行。要把矩陣乘法映射進(jìn)去做法是這樣的把矩陣A和矩陣B分別編碼成紋理。紋理的每個(gè)像素存儲(chǔ)一個(gè)浮點(diǎn)數(shù)或者用RGBA四個(gè)通道存四個(gè)數(shù)。創(chuàng)建一個(gè)和輸出矩陣C同樣大小的幀緩沖作為渲染目標(biāo)。寫一個(gè)片段著色器對于輸出紋理上的每一個(gè)像素位置(i, j)計(jì)算A的第i行和B的第j列的點(diǎn)積。用一次繪制調(diào)用drawArrays或drawElements觸發(fā)整個(gè)計(jì)算。這里有個(gè)細(xì)節(jié)片段著色器不能直接“讀取”另一個(gè)紋理的任意位置嗎可以用texture2D函數(shù)傳入歸一化的紋理坐標(biāo)。但要注意紋理坐標(biāo)是0到1的浮點(diǎn)數(shù)你需要把整數(shù)索引轉(zhuǎn)換成歸一化坐標(biāo)。比如要讀A的第k列坐標(biāo)就是(k 0.5) / width加0.5是為了采樣到像素中心。一次繪制調(diào)用就能算出整個(gè)C矩陣。1000×1000的矩陣乘法在GPU上就是渲染一個(gè)1000×1000的矩形每個(gè)像素做1000次乘加。GPU的幾千個(gè)核心同時(shí)跑實(shí)際耗時(shí)可能在幾毫秒到幾十毫秒之間比CPU快兩個(gè)數(shù)量級。2.3 精度問題WebGL的浮點(diǎn)數(shù)夠用嗎這是很多人會(huì)忽略的坑。WebGL 1.0的片段著色器默認(rèn)使用mediump精度也就是16位浮點(diǎn)數(shù)范圍大約±65504精度只有10位尾數(shù)。用來做神經(jīng)網(wǎng)絡(luò)訓(xùn)練梯度值很容易下溢到0或者上溢到Infinity。訓(xùn)練直接崩掉。解決方案有兩個(gè)第一在著色器里顯式聲明highp精度。precision highp float;這樣可以用到32位浮點(diǎn)數(shù)具體取決于GPU實(shí)現(xiàn)但大多數(shù)現(xiàn)代設(shè)備支持。highp的范圍和精度足夠覆蓋神經(jīng)網(wǎng)絡(luò)訓(xùn)練中的數(shù)值。第二如果設(shè)備不支持highp一些老舊的移動(dòng)GPU可以考慮用“分塊計(jì)算”或者“縮放因子”來規(guī)避。比如把梯度乘以一個(gè)大的常數(shù)計(jì)算完再除回來。但這會(huì)增加復(fù)雜度實(shí)際項(xiàng)目中我建議直接檢測設(shè)備能力不支持highp就降級到CPU路徑。還有一個(gè)隱藏問題WebGL的紋理格式。默認(rèn)的RGBA紋理每個(gè)通道8位只能存0到255的整數(shù)。要存浮點(diǎn)數(shù)必須使用OES_texture_float擴(kuò)展把紋理內(nèi)部格式設(shè)為gl.FLOAT。這個(gè)擴(kuò)展在現(xiàn)代瀏覽器上支持率很高但還是要做特性檢測。3. 把神經(jīng)網(wǎng)絡(luò)拆成WebGL能吃的“圖形任務(wù)”3.1 全連接層的前向傳播一次矩陣乘法加一次激活假設(shè)你有一個(gè)簡單的全連接層輸入向量x維度n權(quán)重矩陣W維度m×n偏置b維度m輸出y activation(Wx b)。在WebGL里這一步可以拆成兩個(gè)渲染過程第一個(gè)過程算z Wx b。把W編碼成紋理x編碼成另一個(gè)紋理或者直接作為uniform數(shù)組傳入如果n不大的話。輸出z是一個(gè)m維向量可以渲染成一個(gè)m×1的紋理。片段著色器里對于輸出紋理的第i個(gè)像素計(jì)算W的第i行和x的點(diǎn)積再加上b[i]。第二個(gè)過程算y activation(z)。這是一個(gè)逐元素操作直接把z紋理作為輸入渲染到y(tǒng)紋理。著色器里對每個(gè)像素應(yīng)用激活函數(shù)比如ReLU就是max(0, z)Sigmoid就是1/(1exp(-z))。這里有個(gè)優(yōu)化點(diǎn)如果網(wǎng)絡(luò)有多層可以把中間結(jié)果一直保留在紋理里不需要讀回CPU。每一層的輸出紋理直接作為下一層的輸入紋理。這樣整個(gè)前向傳播就是一系列繪制調(diào)用數(shù)據(jù)始終在GPU顯存里流轉(zhuǎn)避免了PCIe總線的來回拷貝。3.2 反向傳播的鏈?zhǔn)椒▌t在著色器里怎么表達(dá)反向傳播是訓(xùn)練的核心也是最難映射到WebGL的部分。核心公式是對于每一層先算輸出的梯度dL/dz然后算權(quán)重的梯度dL/dW dL/dz * x^T輸入的梯度dL/dx W^T * dL/dz。在WebGL里dL/dz的計(jì)算通常涉及逐元素乘法比如ReLU的導(dǎo)數(shù)就是z0?1:0和dL/dy逐元素相乘。這個(gè)可以寫成一個(gè)片段著色器輸入是dL/dy紋理和z紋理輸出是dL/dz紋理。dL/dW的計(jì)算是一個(gè)外積dL/dz是m維x是n維外積得到m×n的矩陣。這個(gè)可以映射成一次矩陣乘法把dL/dz當(dāng)作列向量x當(dāng)作行向量。輸出紋理的大小是m×n每個(gè)像素(i,j)的值是dL/dz[i] * x[j]。這個(gè)計(jì)算很簡單但要注意紋理的布局——WebGL紋理是二維的你需要決定怎么把一維向量和二維矩陣對應(yīng)起來。dL/dx的計(jì)算是W^T乘以dL/dz這是一個(gè)矩陣向量乘法。W^T是n×mdL/dz是m維結(jié)果是n維。這個(gè)可以復(fù)用前面矩陣乘法的著色器只是把W的紋理轉(zhuǎn)置一下或者在著色器里交換索引。實(shí)際實(shí)現(xiàn)時(shí)我建議把每一層的反向傳播拆成三個(gè)獨(dú)立的渲染過程算dL/dz、算dL/dW、算dL/dx。每個(gè)過程一個(gè)著色器程序輸入輸出都是紋理。這樣邏輯清晰調(diào)試也方便。3.3 參數(shù)更新在GPU上做還是回CPU做參數(shù)更新是W W - learning_rate * dL/dW。這個(gè)操作本身很簡單逐元素減法。但問題在于更新后的W需要作為下一輪前向傳播的輸入。如果W一直留在GPU紋理里那更新也可以在GPU上做寫一個(gè)片段著色器輸入W紋理和dL/dW紋理輸出新的W紋理。但這里有個(gè)權(quán)衡。如果所有參數(shù)都在GPU上你就沒法方便地查看訓(xùn)練過程中的權(quán)重變化也沒法做復(fù)雜的優(yōu)化器比如Adam需要維護(hù)一階矩和二階矩計(jì)算量不大但邏輯稍復(fù)雜。我的做法是前向和反向在GPU上跑參數(shù)更新回CPU做。每一輪迭代結(jié)束后把dL/dW從GPU讀回CPU用readPixels在JavaScript里更新權(quán)重再把新權(quán)重上傳到GPU紋理。這樣雖然多了兩次數(shù)據(jù)傳輸?shù)`活性大大提高而且對于小網(wǎng)絡(luò)來說傳輸開銷遠(yuǎn)小于計(jì)算開銷。讀回?cái)?shù)據(jù)時(shí)要注意readPixels是同步操作會(huì)阻塞渲染管線。如果每輪都讀可能會(huì)成為瓶頸。優(yōu)化方法是批量讀——積累幾輪的梯度再一起讀回或者用WebGL 2.0的PBO像素緩沖對象做異步讀回。但WebGL 2.0的兼容性還不如1.0實(shí)際項(xiàng)目中要權(quán)衡。4. 手寫一個(gè)WebGL神經(jīng)網(wǎng)絡(luò)訓(xùn)練器的完整路徑4.1 環(huán)境搭建與基礎(chǔ)工具函數(shù)首先你需要一個(gè)WebGL上下文。創(chuàng)建一個(gè)canvas獲取gl對象。然后檢查必要的擴(kuò)展const canvas document.createElement(canvas); const gl canvas.getContext(webgl); const floatExt gl.getExtension(OES_texture_float); if (!floatExt) { throw new Error(設(shè)備不支持浮點(diǎn)紋理無法進(jìn)行GPU訓(xùn)練); }接下來是幾個(gè)基礎(chǔ)工具函數(shù)這些函數(shù)會(huì)反復(fù)用到創(chuàng)建浮點(diǎn)紋理輸入是一個(gè)Float32Array輸出是一個(gè)WebGL紋理對象。注意設(shè)置紋理參數(shù)為NEAREST采樣避免插值導(dǎo)致數(shù)值錯(cuò)誤。function createFloatTexture(gl, data, width, height) { const texture gl.createTexture(); gl.bindTexture(gl.TEXTURE_2D, texture); gl.texImage2D(gl.TEXTURE_2D, 0, gl.RGBA, width, height, 0, gl.RGBA, gl.FLOAT, data); gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MIN_FILTER, gl.NEAREST); gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MAG_FILTER, gl.NEAREST); gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_S, gl.CLAMP_TO_EDGE); gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_T, gl.CLAMP_TO_EDGE); return texture; }創(chuàng)建幀緩沖幀緩沖是渲染目標(biāo)你把片段著色器的輸出寫進(jìn)去。每個(gè)幀緩沖綁定一個(gè)紋理作為顏色附件。function createFramebuffer(gl, texture) { const fb gl.createFramebuffer(); gl.bindFramebuffer(gl.FRAMEBUFFER, fb); gl.framebufferTexture2D(gl.FRAMEBUFFER, gl.COLOR_ATTACHMENT0, gl.TEXTURE_2D, texture, 0); return fb; }編譯著色器程序這是標(biāo)準(zhǔn)流程編譯頂點(diǎn)著色器和片段著色器鏈接成程序。頂點(diǎn)著色器通常很簡單就是傳遞一個(gè)覆蓋整個(gè)畫布的矩形。function createProgram(gl, vertexSource, fragmentSource) { const vs gl.createShader(gl.VERTEX_SHADER); gl.shaderSource(vs, vertexSource); gl.compileShader(vs); const fs gl.createShader(gl.FRAGMENT_SHADER); gl.shaderSource(fs, fragmentSource); gl.compileShader(fs); const program gl.createProgram(); gl.attachShader(program, vs); gl.attachShader(program, fs); gl.linkProgram(program); return program; }這些工具函數(shù)看起來簡單但每個(gè)都有坑。比如紋理的寬高必須是2的冪嗎在WebGL 1.0里如果使用NEAREST采樣和CLAMP_TO_EDGE可以不是2的冪。但有些驅(qū)動(dòng)會(huì)強(qiáng)制要求所以最安全的做法是padding到2的冪。我一般會(huì)把矩陣維度補(bǔ)齊到最近的2的冪比如1000×1000的矩陣實(shí)際分配1024×1024的紋理多出來的部分填0。4.2 矩陣乘法著色器的逐行拆解矩陣乘法是整個(gè)訓(xùn)練流程里最核心的計(jì)算。假設(shè)我們要算C A × BA是M×KB是K×NC是M×N。頂點(diǎn)著色器很簡單就是畫一個(gè)覆蓋整個(gè)輸出紋理的矩形attribute vec2 a_position; void main() { gl_Position vec4(a_position, 0.0, 1.0); }片段著色器是重點(diǎn)precision highp float; uniform sampler2D u_A; uniform sampler2D u_B; uniform vec2 u_A_size; // (M, K) uniform vec2 u_B_size; // (K, N) uniform vec2 u_C_size; // (M, N) void main() { vec2 c_pos gl_FragCoord.xy; // 當(dāng)前像素在C中的位置 float sum 0.0; for (int k 0; k int(u_A_size.y); k) { // 讀A的第c_pos.x行第k列 vec2 a_uv vec2((c_pos.x 0.5) / u_A_size.x, (float(k) 0.5) / u_A_size.y); float a_val texture2D(u_A, a_uv).r; // 讀B的第k行第c_pos.y列 vec2 b_uv vec2((float(k) 0.5) / u_B_size.x, (c_pos.y 0.5) / u_B_size.y); float b_val texture2D(u_B, b_uv).r; sum a_val * b_val; } gl_FragColor vec4(sum, 0.0, 0.0, 1.0); }這段代碼有幾個(gè)關(guān)鍵點(diǎn)第一gl_FragCoord.xy給出的是像素中心坐標(biāo)比如第一個(gè)像素是(0.5, 0.5)。所以計(jì)算紋理坐標(biāo)時(shí)要加0.5再除以尺寸。第二循環(huán)次數(shù)必須是常量或者可以用uniform控制。WebGL 1.0的GLSL不允許非常量循環(huán)邊界但可以用一個(gè)足夠大的常量循環(huán)然后在循環(huán)體內(nèi)用if判斷是否超出范圍。不過現(xiàn)代瀏覽器對動(dòng)態(tài)循環(huán)的支持已經(jīng)不錯(cuò)了實(shí)際測試中大多數(shù)設(shè)備都能跑。第三每次循環(huán)要讀兩次紋理。紋理讀取是GPU計(jì)算中最慢的操作之一因?yàn)橐L問顯存。優(yōu)化方法是把A和B的紋理綁定到不同的紋理單元并且盡量讓循環(huán)內(nèi)的訪問模式有局部性。但矩陣乘法的訪問模式天然是跳躍的所以這個(gè)優(yōu)化空間有限。一個(gè)實(shí)用的優(yōu)化技巧分塊計(jì)算。把大矩陣切成小塊比如32×32的塊每個(gè)塊的計(jì)算可以更好地利用GPU的緩存。但這會(huì)增加著色器的復(fù)雜度實(shí)際項(xiàng)目中如果矩陣不大比如全連接層通常幾百維直接算就夠了。4.3 激活函數(shù)與損失函數(shù)的GPU實(shí)現(xiàn)激活函數(shù)是逐元素操作實(shí)現(xiàn)起來簡單。以ReLU為例precision highp float; uniform sampler2D u_input; void main() { vec2 uv gl_FragCoord.xy / vec2(textureSize(u_input, 0)); float val texture2D(u_input, uv).r; gl_FragColor vec4(max(0.0, val), 0.0, 0.0, 1.0); }Sigmoid稍微復(fù)雜一點(diǎn)因?yàn)橐鉫xpfloat sigmoid(float x) { return 1.0 / (1.0 exp(-x)); }但要注意exp函數(shù)在highp下可能會(huì)有精度問題特別是當(dāng)x很大或很小時(shí)。一個(gè)穩(wěn)定的實(shí)現(xiàn)是float stable_sigmoid(float x) { if (x 0.0) { return 1.0 / (1.0 exp(-x)); } else { float e exp(x); return e / (1.0 e); } }損失函數(shù)通常用交叉熵。對于分類問題交叉熵的梯度是softmax輸出減去真實(shí)標(biāo)簽。這個(gè)計(jì)算可以放在CPU做因?yàn)橹簧婕拜敵鰧佑?jì)算量小。但如果想全GPU化也可以寫一個(gè)著色器輸入是預(yù)測概率和真實(shí)標(biāo)簽輸出是每個(gè)樣本的損失值。4.4 反向傳播的著色器實(shí)現(xiàn)細(xì)節(jié)反向傳播里最復(fù)雜的是dL/dW的計(jì)算。假設(shè)當(dāng)前層的輸入是xn維輸出梯度是dL/dzm維那么dL/dW是一個(gè)m×n的矩陣每個(gè)元素是dL/dz[i] * x[j]。這個(gè)計(jì)算可以寫成一個(gè)著色器輸出紋理大小是m×nprecision highp float; uniform sampler2D u_dz; // m維向量存成m×1紋理 uniform sampler2D u_x; // n維向量存成n×1紋理 uniform vec2 u_dz_size; uniform vec2 u_x_size; void main() { vec2 pos gl_FragCoord.xy; // pos.x是行索引ipos.y是列索引j vec2 dz_uv vec2(0.5 / u_dz_size.x, (pos.x 0.5) / u_dz_size.y); vec2 x_uv vec2(0.5 / u_x_size.x, (pos.y 0.5) / u_x_size.y); float dz_val texture2D(u_dz, dz_uv).r; float x_val texture2D(u_x, x_uv).r; gl_FragColor vec4(dz_val * x_val, 0.0, 0.0, 1.0); }注意這里紋理坐標(biāo)的寫法。因?yàn)閐z和x都是一維向量我習(xí)慣把它們存成寬度為1、高度為維度的紋理。這樣讀取時(shí)x坐標(biāo)固定為0.5/width即像素中心y坐標(biāo)根據(jù)索引變化。dL/dx的計(jì)算是W^T × dL/dz。W是m×nW^T是n×mdL/dz是m維結(jié)果是n維。這個(gè)可以復(fù)用矩陣乘法的著色器只需要把W的紋理坐標(biāo)交換一下?；蛘邔懸粋€(gè)專門的著色器precision highp float; uniform sampler2D u_W; // m×n uniform sampler2D u_dz; // m維 uniform vec2 u_W_size; uniform vec2 u_dz_size; void main() { float i gl_FragCoord.y; // 輸出向量的索引 float sum 0.0; for (int j 0; j int(u_W_size.x); j) { vec2 w_uv vec2((float(j) 0.5) / u_W_size.x, (i 0.5) / u_W_size.y); float w_val texture2D(u_W, w_uv).r; vec2 dz_uv vec2(0.5 / u_dz_size.x, (float(j) 0.5) / u_dz_size.y); float dz_val texture2D(u_dz, dz_uv).r; sum w_val * dz_val; } gl_FragColor vec4(sum, 0.0, 0.0, 1.0); }這里W的讀取是W[j][i]因?yàn)槲覀円鉝^T的第i行也就是W的第i列。所以w_uv的x坐標(biāo)是j行索引y坐標(biāo)是i列索引。這個(gè)索引交換很容易搞混寫的時(shí)候要特別小心。5. 實(shí)測中那些文檔不會(huì)告訴你的坑5.1 紋理尺寸限制與矩陣分塊策略WebGL有一個(gè)硬性限制MAX_TEXTURE_SIZE。這個(gè)值在不同設(shè)備上不一樣低端手機(jī)可能只有2048高端顯卡能到16384。如果你的權(quán)重矩陣超過這個(gè)尺寸就必須分塊。分塊矩陣乘法的思路是把A切成若干行塊B切成若干列塊然后逐塊計(jì)算。每個(gè)塊的計(jì)算結(jié)果寫到一個(gè)臨時(shí)紋理最后拼起來。但這樣會(huì)增加繪制調(diào)用次數(shù)和紋理切換開銷。我的經(jīng)驗(yàn)是對于全連接層權(quán)重矩陣通常不會(huì)超過4096×4096那已經(jīng)是1600萬參數(shù)了所以大多數(shù)設(shè)備都能直接放下。但如果做卷積層特征圖可能很大就需要分塊。實(shí)際項(xiàng)目中我建議先檢測MAX_TEXTURE_SIZE如果矩陣維度超過限制就自動(dòng)降級到CPU計(jì)算或者用分塊策略。還有一個(gè)隱藏限制MAX_RENDERBUFFER_SIZE和MAX_VIEWPORT_DIMS。這些通常和紋理尺寸限制一致但有些設(shè)備上幀緩沖的尺寸限制更嚴(yán)格。創(chuàng)建幀緩沖后一定要檢查checkFramebufferStatus確保返回FRAMEBUFFER_COMPLETE。5.2 浮點(diǎn)紋理的讀寫精度陷阱前面提到要用OES_texture_float擴(kuò)展。但即使擴(kuò)展存在不同設(shè)備對浮點(diǎn)紋理的支持程度也不一樣。有些設(shè)備只能渲染到浮點(diǎn)紋理但不能讀取有些設(shè)備讀取時(shí)會(huì)有精度損失。一個(gè)實(shí)用的檢測方法是創(chuàng)建一個(gè)浮點(diǎn)紋理寫入一個(gè)已知值比如0.123456789然后讀回來看精度是否滿足要求。如果誤差超過1e-6就說明這個(gè)設(shè)備的浮點(diǎn)紋理不可靠需要降級。另外readPixels讀回浮點(diǎn)數(shù)據(jù)時(shí)返回的是Float32Array。但有些設(shè)備會(huì)返回Uint8Array即使你請求的是浮點(diǎn)格式。所以讀回后要檢查數(shù)據(jù)類型必要時(shí)做轉(zhuǎn)換。還有一個(gè)坑WebGL的texImage2D在傳入Float32Array時(shí)如果數(shù)據(jù)長度和width×height×4不匹配會(huì)靜默失敗或者產(chǎn)生未定義行為。一定要確保數(shù)據(jù)長度正確。我習(xí)慣在創(chuàng)建紋理前先算好width * height * 4然后檢查數(shù)組長度。5.3 訓(xùn)練不收斂的排查鏈路GPU訓(xùn)練最讓人頭疼的問題是訓(xùn)練不收斂但你看不到中間過程。CPU訓(xùn)練時(shí)你可以打印每一層的梯度、權(quán)重但GPU上數(shù)據(jù)都在紋理里讀回來很麻煩。我總結(jié)了一套排查鏈路第一步檢查前向傳播。用一個(gè)已知的輸入手動(dòng)算一遍CPU版本的前向結(jié)果然后和GPU版本對比。如果前向就不對那反向肯定不對。對比時(shí)要注意精度GPU的浮點(diǎn)運(yùn)算和CPU可能有微小差異但應(yīng)該在1e-5以內(nèi)。第二步檢查梯度。用數(shù)值梯度法驗(yàn)證對于某個(gè)權(quán)重w計(jì)算Loss(wε)和Loss(w-ε)然后(Loss(wε)-Loss(w-ε))/(2ε)就是數(shù)值梯度。和GPU算出的梯度對比如果差異大說明反向傳播有問題。第三步檢查參數(shù)更新。如果梯度對了但Loss不下降可能是學(xué)習(xí)率太大或太小。GPU訓(xùn)練時(shí)梯度值可能因?yàn)榫葐栴}被截?cái)鄬?dǎo)致實(shí)際學(xué)習(xí)率和設(shè)定值不符?？梢韵劝褜W(xué)習(xí)率調(diào)小10倍試試。第四步檢查數(shù)據(jù)。有時(shí)候問題不在GPU而在輸入數(shù)據(jù)。比如標(biāo)簽編碼錯(cuò)了或者輸入沒有歸一化。GPU訓(xùn)練對輸入尺度很敏感因?yàn)楦↑c(diǎn)數(shù)的動(dòng)態(tài)范圍有限。5.4 性能調(diào)優(yōu)從每秒幾輪到每秒幾百輪最初的實(shí)現(xiàn)可能每秒只能跑幾輪迭代慢得讓人懷疑人生。優(yōu)化方向有幾個(gè)減少繪制調(diào)用。每次繪制都有固定開銷包括狀態(tài)切換、著色器綁定、紋理綁定。如果能把多個(gè)操作合并到一個(gè)著色器里就能減少調(diào)用次數(shù)。比如把矩陣乘法和偏置加法合并把激活函數(shù)和損失函數(shù)合并。使用頂點(diǎn)緩沖對象VBO。畫矩形時(shí)不要每次都用gl.vertexAttribPointer重新指定頂點(diǎn)數(shù)據(jù)。創(chuàng)建一個(gè)VBO把矩形的四個(gè)頂點(diǎn)存進(jìn)去每次繪制時(shí)直接綁定。避免不必要的紋理切換。紋理切換是昂貴的操作。如果多個(gè)著色器使用相同的輸入紋理盡量把它們安排在連續(xù)的繪制調(diào)用里。使用WebGL 2.0。WebGL 2.0支持更多紋理格式、多重渲染目標(biāo)MRT、以及更靈活的循環(huán)控制。如果目標(biāo)用戶設(shè)備支持優(yōu)先用WebGL 2.0。但要注意兼容性做好降級方案。批處理。如果訓(xùn)練數(shù)據(jù)是小樣本可以把多個(gè)樣本打包成一個(gè)批次一次前向傳播算完。這樣GPU的利用率更高因?yàn)槊總€(gè)著色器調(diào)用處理的數(shù)據(jù)量更大。實(shí)測下來一個(gè)兩層全連接網(wǎng)絡(luò)784-128-10在MNIST上訓(xùn)練優(yōu)化前每秒約5輪優(yōu)化后能到每秒200輪以上。這個(gè)速度已經(jīng)足夠做交互式教學(xué)了——學(xué)生寫一個(gè)數(shù)字模型在幾百毫秒內(nèi)就能完成一輪微調(diào)。6. 瀏覽器端訓(xùn)練的邊界與適用場景6.1 什么規(guī)模的模型適合放在瀏覽器里訓(xùn)不是所有模型都適合瀏覽器端訓(xùn)練。我的經(jīng)驗(yàn)是參數(shù)量在10萬到100萬之間的模型WebGL訓(xùn)練是可行的。超過這個(gè)規(guī)模紋理尺寸和顯存會(huì)成為瓶頸而且訓(xùn)練時(shí)間會(huì)變得不可接受。具體來說全連接層每層幾百到幾千維層數(shù)不超過5層這樣的網(wǎng)絡(luò)在瀏覽器里訓(xùn)練很流暢。卷積層也可以做但需要把卷積展開成矩陣乘法im2col這會(huì)增加內(nèi)存占用。如果卷積核不大比如3×3通道數(shù)不多比如32或64也是可以跑的。不適合的場景包括大型語言模型、高分辨率圖像生成模型、需要大量迭代才能收斂的模型。這些還是交給專業(yè)訓(xùn)練環(huán)境。6.2 隱私保護(hù)與離線可用性的真實(shí)價(jià)值瀏覽器端訓(xùn)練最大的賣點(diǎn)不是性能而是數(shù)據(jù)不出本地。用戶的輸入數(shù)據(jù)、標(biāo)注數(shù)據(jù)、甚至模型本身都可以完全留在瀏覽器里。這對于隱私敏感的應(yīng)用場景很有價(jià)值比如醫(yī)療健康類的個(gè)人數(shù)據(jù)訓(xùn)練、教育類的學(xué)生行為數(shù)據(jù)建模。另一個(gè)價(jià)值是離線可用。一旦頁面加載完成后續(xù)的訓(xùn)練和推理都不需要網(wǎng)絡(luò)連接。這對于網(wǎng)絡(luò)不穩(wěn)定的環(huán)境或者需要完全離線部署的場景是很實(shí)用的特性。還有一個(gè)容易被忽略的價(jià)值即時(shí)反饋。用戶補(bǔ)充幾個(gè)樣本模型立刻更新識別效果當(dāng)場變好。這種交互體驗(yàn)是云端訓(xùn)練做不到的因?yàn)樵贫擞?xùn)練有網(wǎng)絡(luò)延遲和排隊(duì)等待。6.3 和TensorFlow.js的GPU后端有什么不同TensorFlow.js也支持WebGL后端而且封裝得很好幾行代碼就能定義模型并訓(xùn)練。那為什么還要自己寫WebGL著色器區(qū)別在于控制粒度。TensorFlow.js的WebGL后端是一個(gè)通用框架它會(huì)把各種操作映射到著色器但中間有很多抽象層。對于標(biāo)準(zhǔn)操作比如全連接、卷積、常見的激活函數(shù)它的性能已經(jīng)很好。但如果你需要自定義操作或者想針對特定硬件做極致優(yōu)化自己寫著色器更靈活。另外TensorFlow.js的WebGL后端在某些設(shè)備上會(huì)有兼容性問題而且它的內(nèi)存管理是自動(dòng)的有時(shí)候會(huì)導(dǎo)致紋理泄漏或者頻繁的垃圾回收。自己管理紋理和幀緩沖雖然代碼量大但行為完全可控。我的建議是如果標(biāo)準(zhǔn)模型能滿足需求直接用TensorFlow.js省時(shí)省力。如果需要自定義訓(xùn)練邏輯或者對性能有極致要求再考慮手寫WebGL。7. 從原型到產(chǎn)品工程化落地的幾個(gè)關(guān)鍵決策7.1 如何設(shè)計(jì)降級方案保證兼容性不是所有瀏覽器都支持WebGL也不是所有支持WebGL的設(shè)備都支持浮點(diǎn)紋理。產(chǎn)品化時(shí)必須設(shè)計(jì)降級方案。我的做法是三級降級第一級WebGL 2.0 浮點(diǎn)紋理。這是最優(yōu)路徑性能最好。第二級WebGL 1.0 浮點(diǎn)紋理擴(kuò)展。性能稍差但兼容性更廣。第三級純CPU JavaScript。用TypedArray做矩陣運(yùn)算雖然慢但保證能用。對于小模型參數(shù)量小于1萬CPU版本也能在可接受的時(shí)間內(nèi)完成訓(xùn)練。降級檢測在頁面加載時(shí)做一次把結(jié)果緩存起來。用戶不需要知道底層用的是哪條路徑只需要看到訓(xùn)練能跑起來。7.2 訓(xùn)練過程的進(jìn)度反饋與中斷處理GPU訓(xùn)練是異步的你發(fā)出繪制調(diào)用后GPU在后臺跑JavaScript繼續(xù)執(zhí)行。如果訓(xùn)練循環(huán)寫成一個(gè)大的for循環(huán)頁面會(huì)卡死用戶看不到任何進(jìn)度。解決方案是用requestAnimationFrame或者setTimeout把訓(xùn)練循環(huán)拆成小批次。每一幀跑幾輪迭代然后更新進(jìn)度條再請求下一幀。這樣頁面保持響應(yīng)用戶能看到訓(xùn)練進(jìn)度。中斷處理也很重要。用戶可能中途關(guān)閉頁面或者切換標(biāo)簽頁。要在beforeunload事件里清理WebGL資源避免內(nèi)存泄漏。另外如果用戶切換標(biāo)簽頁requestAnimationFrame會(huì)暫停訓(xùn)練也會(huì)暫停。這是符合預(yù)期的行為但要在UI上給出提示。7.3 模型持久化把訓(xùn)練好的權(quán)重存到哪里訓(xùn)練完成后權(quán)重需要保存。GPU紋理里的數(shù)據(jù)不能直接序列化必須先讀回CPU轉(zhuǎn)成Float32Array然后存到IndexedDB或者localStorage。IndexedDB適合存大文件幾MB到幾十MBlocalStorage只適合存小數(shù)據(jù)幾KB。對于神經(jīng)網(wǎng)絡(luò)權(quán)重通常用IndexedDB。存儲(chǔ)時(shí)可以把Float32Array轉(zhuǎn)成Blob或者直接存ArrayBuffer。讀取時(shí)反向操作從IndexedDB讀出ArrayBuffer轉(zhuǎn)成Float32Array再上傳到GPU紋理。整個(gè)過程要注意異步處理因?yàn)镮ndexedDB的讀寫都是異步的。還有一個(gè)細(xì)節(jié)權(quán)重的版本管理。如果模型結(jié)構(gòu)變了舊權(quán)重就不能用了。我習(xí)慣在存儲(chǔ)時(shí)帶上一個(gè)版本號或者結(jié)構(gòu)哈希讀取時(shí)校驗(yàn)不匹配就重新初始化。8. 我踩過的三個(gè)印象最深的坑第一個(gè)坑是紋理坐標(biāo)的0.5偏移。剛開始寫著色器時(shí)我直接用了gl_FragCoord.xy / textureSize作為紋理坐標(biāo)結(jié)果采樣到的總是相鄰像素的值導(dǎo)致矩陣乘法結(jié)果完全錯(cuò)誤。排查了很久才發(fā)現(xiàn)gl_FragCoord給出的是像素中心坐標(biāo)而紋理坐標(biāo)需要的是歸一化的像素中心所以要加0.5再除。這個(gè)細(xì)節(jié)在大多數(shù)WebGL教程里都不會(huì)強(qiáng)調(diào)但做通用計(jì)算時(shí)是致命的。第二個(gè)坑是浮點(diǎn)紋理的線性插值。默認(rèn)情況下紋理采樣使用線性插值會(huì)混合相鄰像素的值。對于圖像渲染這是好事但對于矩陣運(yùn)算這是災(zāi)難。必須把TEXTURE_MIN_FILTER和TEXTURE_MAG_FILTER都設(shè)為NEAREST。我一開始忘了設(shè)結(jié)果梯度計(jì)算出來總是有微小的噪聲訓(xùn)練完全不收斂。第三個(gè)坑是幀緩沖的完整性檢查。創(chuàng)建幀緩沖后我沒有檢查checkFramebufferStatus直接就開始渲染。在某些設(shè)備上幀緩沖創(chuàng)建失敗但不會(huì)報(bào)錯(cuò)渲染結(jié)果全是0。后來養(yǎng)成了習(xí)慣每次創(chuàng)建幀緩沖后都檢查狀態(tài)不完整就拋異常。這個(gè)檢查在開發(fā)階段能省下大量調(diào)試時(shí)間。9. 后續(xù)可以繼續(xù)深挖的方向如果你已經(jīng)把基礎(chǔ)的WebGL訓(xùn)練跑通了還有幾個(gè)方向可以繼續(xù)優(yōu)化。一是用WebGL 2.0的變換反饋Transform Feedback。WebGL 2.0允許把頂點(diǎn)著色器的輸出直接寫回緩沖區(qū)不需要經(jīng)過片段著色器。這對于逐元素操作比如激活函數(shù)、參數(shù)更新可能更高效因?yàn)楸苊饬思y理的創(chuàng)建和綁定。二是混合精度訓(xùn)練。用16位浮點(diǎn)數(shù)做前向和反向用32位浮點(diǎn)數(shù)做參數(shù)更新。這樣可以減少顯存占用和帶寬消耗速度可能提升一倍。但要注意數(shù)值穩(wěn)定性可能需要損失縮放loss scaling來防止梯度下溢。三是多GPU協(xié)同。如果用戶設(shè)備有多個(gè)GPU比如集成顯卡加獨(dú)立顯卡理論上可以用多個(gè)WebGL上下文并行計(jì)算。但瀏覽器對多GPU的支持還很有限實(shí)際可行性不高。四是和WebAssembly結(jié)合。把一些不適合GPU的操作比如數(shù)據(jù)預(yù)處理、復(fù)雜的優(yōu)化器邏輯放到WASM里做GPU專注做矩陣運(yùn)算。兩者通過SharedArrayBuffer通信實(shí)現(xiàn)真正的并行。這條路我走了大半年從最初的“這能行嗎”到現(xiàn)在的“基本能用”中間踩了無數(shù)坑但也積累了不少經(jīng)驗(yàn)。瀏覽器端訓(xùn)練不是要取代專業(yè)訓(xùn)練環(huán)境而是在特定場景下提供一種輕量、隱私、即時(shí)的選擇。如果你的項(xiàng)目恰好需要這些特性WebGL方案值得一試。