TensorFlow.js 实战笔记：浏览器端序列模型训练的内存、后端与性能陷阱

写在前面

过去几个月我一直在做一个浏览器端的机器学习实验工具，用 TensorFlow.js 在用户本地训练 LSTM 和 Transformer 模型。整个过程中踩了不少坑——有些是文档里没写的，有些是官方示例不会遇到的。这篇文章把这些经验整理出来，希望能帮到同样在浏览器端跑模型的同学。

文章不聊概念，只聊实际遇到的问题和解决方式。

一、后端选型：别默认用 WebGL

TensorFlow.js 支持 CPU、WebGL、WASM、WebGPU 四种后端。大多数人默认选 WebGL，因为它“看起来最快”。但实际用下来，WebGL 有两个致命问题：

第一是冷启动慢。 首次推理时，着色器编译可能要几百毫秒甚至更久。用户打开页面第一次点“训练”，会明显卡一下。我的解决方式是在应用初始化时，用一个小张量跑一次推理做预热：

// 页面加载时预热 WebGL 后端
await tf.setBackend('webgl');
await tf.ready();
const warmup = tf.tensor1d([1, 2, 3]);
warmup.square().dataSync();
warmup.dispose();

第二是设备兼容性差。 不同 GPU、不同浏览器对 WebGL 的支持差异很大。我在一台老 MacBook 上测试时，同样的模型在 Chrome 里能跑，在 Safari 里直接报 “shader compile failed”。

我的最终选择是 WASM。 它比 WebGL 慢一些，但稳定性好得多，跨设备表现一致。对于序列模型这种计算量不算特别大的场景，稳定性比极致性能更重要。

如果你的模型以矩阵乘为主，且用户主要在 Chrome 113+ 上访问，可以试试 WebGPU。但不要把它当作主力后端，兼容性还不到时候。

二、内存管理：tf.tidy 不是万能药

浏览器端跑模型最大的坑是内存泄漏。JS 的 GC 管不了 WebGL 显存，必须手动释放张量。TF.js 提供了 tf.tidy() 来自动清理，但它有几个限制，文档里写得不明显。

限制一：tf.tidy 只处理同步函数。

下面这样写是无效的：

// ❌ 错误：tidy 不接受 async 函数
const result = await tf.tidy(async () => {
    const a = tf.tensor([1, 2, 3]);
    return await someAsyncOp(a);
});

异步场景下，需要用 tf.engine().startScope() 和 endScope() 手动管理：

// ✅ 正确：手动管理作用域
async function predict(input) {
    tf.engine().startScope();
    try {
        const output = model.predict(input);
        return await output.data();
    } finally {
        tf.engine().endScope();
    }
}

限制二：tf.tidy 不清理变量。

tf.variable() 创建的变量是持久的，即使写在 tidy 里也不会被自动清理。释放变量必须显式调用 variable.dispose() 或 tf.disposeVariables()。

限制三：返回值的“传染性”。

tidy 返回的张量不会被清理，但它引用的中间张量会被清理。如果你的返回值是某个中间张量的 view（比如 slice 的结果），要小心它是否还依赖被释放的内存。

我的做法是在训练循环里监控张量数量：

await model.fit(xs, ys, {
    epochs,
    callbacks: {
        onEpochEnd: (epoch, logs) => {
            const numTensors = tf.memory().numTensors;
            console.log(`Epoch ${epoch}: loss=${logs.loss}, tensors=${numTensors}`);
            // 正常情况下这个数字应该稳定波动，持续增长就是泄漏
            if (numTensors > 500) {
                console.warn('Tensor count too high, possible leak');
            }
        },
    },
});

这个简单的日志帮我定位过好几次泄漏。有一次是数据预处理时忘了 tidy，每轮泄漏几十个张量，跑十几轮就崩了。

三、训练性能：数据在 CPU 和 GPU 之间来回搬是最大的浪费

TF.js 的训练流程里，数据在 CPU 和 GPU 之间来回搬运是最耗时的环节。典型流程是：

1. 数据在 CPU 内存里
2. 上传到 GPU 显存（慢）
3. GPU 上计算
4. 结果读回 CPU（慢）

每次 readback 都会让 GPU 流水线停顿。TF.js 3.13.0 引入了 tensor.dataToGPU()，可以直接拿到 GPU 资源，让下游处理继续在 GPU 上跑，不用落回 CPU：

const gpuData = tensor.dataToGPU({ customTexShape: [128, 128] });
// WebGL: gpuData.texture
// WebGPU: gpuData.buffer

但大部分应用场景用不上这么底层的 API。 对我们来说，更实际的优化是：

• 减少 readback 频率：不要每个 batch 都 .dataSync() 读回损失值，让它留在 GPU 上累积。
• 用 tf.batch() 合并小张量：小张量频繁创建会触发大量 GPU 调用，合并后能明显提速。
• 避免在训练循环里做 CPU 密集操作：比如每轮都重新 Array.from() 转换数据，会让 GPU 空转等 CPU。

四、自定义 Kernel 和梯度：CRF 层的实现经验

TF.js 没有内置 CRF 层，但序列标注任务又绕不开它。我参考官方文档实现了自定义 Kernel 和 Gradient。

Kernel 的注册：

import { registerKernel } from '@tensorflow/tfjs-core';

registerKernel({
    kernelName: 'CustomCRFDecode',
    backendName: 'wasm',
    kernelFunc: ({ inputs, attrs }) => {
        const { emissions } = inputs;
        // Viterbi 解码的实现
        return decode(emissions, attrs.transitions);
    },
});

梯度的注册有两种方式：

第一种是通过 registerGradient，注册到 Gradient Registry：

import { registerGradient } from '@tensorflow/tfjs-core';

registerGradient('CustomCRFDecode', {
    'emissions': (dy, y, emissions, transitions) => {
        // 返回对 emissions 的梯度
    },
});

第二种是 tf.customGrad，它绕过 Gradient Registry，直接在函数内部定义梯度计算逻辑，更灵活，适合研究场景。

我踩的坑是： 自定义 Kernel 是后端特定的，WebGL 和 WASM 需要各写一份实现。如果只在 WASM 上注册了 Kernel，切到 WebGL 后端会直接报 “kernel not found”。对于只需要支持一个后端的项目，这不算问题；但对需要多后端切换的项目，工作量翻倍。

最终我放弃了完整的 CRF 实现，改用了一个更轻量的方案：在 softmax 输出后加一层掩码（mask），把已经被选过的号码概率置零，再重新归一化。效果上接近 CRF 的约束，但实现成本低得多。

五、模型加载：别把所有权重都塞进 bundle

如果你的模型是从服务器加载的，不要把 .bin 权重文件和 JS 代码打包在一起。权重文件通常有几 MB 到几十 MB，会显著拖慢首屏。

正确的做法是把权重文件放在 CDN 或对象存储上，用 tf.loadLayersModel() 按需加载：

const model = await tf.loadLayersModel('/models/ssq-lstm/model.json');

TF.js 会自动请求 model.json 里引用的权重分片文件。这些文件走 HTTP 缓存，用户第二次访问就不用重新下载了。

对于需要用户本地训练的场景，可以在模型结构定义好之后，用 model.setWeights() 从预训练的权重初始化，或者干脆从随机初始化开始。

六、一个实用的调试技巧

浏览器端训练没有 Python 那样成熟的调试工具，我的做法是在关键节点插入“张量快照”：

function inspectTensor(name, tensor) {
    const data = tensor.dataSync();
    console.log(`[${name}] shape=${tensor.shape}, min=${Math.min(...data)}, max=${Math.max(...data)}, mean=${data.reduce((a,b)=>a+b)/data.length}`);
}

在模型 build 完成、每个 epoch 结束时打印关键张量的形状和统计信息，能快速发现“梯度爆炸”或“输出全为 0”这类问题。

七、总结

浏览器端跑机器学习，最大的挑战不是模型本身，而是资源受限环境下的工程细节。几个核心经验：

1. 后端选型优先考虑稳定性，WASM 比 WebGL 更可控
2. 内存管理必须显式处理，tidy 不是万能的，异步场景要手动管理作用域
3. 减少 CPU-GPU 数据搬运，是性能优化的关键
4. 自定义 Kernel 是后端特定的，多后端项目要评估工作量
5. 模型权重走 CDN，不要塞进 bundle

TensorFlow.js 的文档偏概念，很多实战细节要靠踩坑积累。希望这篇文章能帮你少走一些弯路。