在移动设备上运行Transformer类模型长期被认为是不可能的任务，直到MobileBERT这类精简架构出现。MobileBERT通过保持层数接近BERT_base但大幅缩减隐层维度，配合瓶颈结构，在精度损失很小的前提下把计算量降了下来。Node.js作为服务端和跨端脚本运行时，借助TensorFlow.js可以在不写原生代码的情况下加载并推理MobileBERT，这在React Native或Electron打包的App里尤其有用。本文从工程落地角度，说明如何把模型跑通、如何控制资源、如何处理中文输入。
模型转换与TensorFlow.js加载
MobileBERT原始权重多为TensorFlow SavedModel或PyTorch格式，直接给Node.js用并不现实。需要先转成TensorFlow.js的GraphModel格式，也就是model.json加分片二进制权重。转换时使用tensorflowjs_converter并开启量化，能把浮点32位压到16位甚至8位，显著缩小包体。注意MobileBERT里有些自定义算子，在标准转换脚本中可能报错，需要先在Python侧把模型导出为精简计算图，再转换。
在Node.js侧，我们用@tensorflow/tfjs-node而不是纯JS版，因为前者调用了本地二进制，推理速度能提升数倍。加载模型只需tf.loadGraphModel，但要注意移动端文件系统路径处理。下面代码展示了一个最小加载示例，其中包含对模型地址的错误处理，避免App白屏。
const tf = require('@tensorflow/tfjs-node');
const path = require('path');
async function loadMobileBERT(modelDir) {
const modelPath = 'file://' + path.join(modelDir, 'model.json');
try {
const model = await tf.loadGraphModel(modelPath);
// 预热一次，避免首帧延迟过高
const dummy = tf.zeros([1, 128], 'int32');
await model.executeAsync({ input_ids: dummy, attention_mask: dummy });
dummy.dispose();
return model;
} catch (e) {
console.error('模型加载失败:', e.message);
return null;
}
}
量化虽然减小体积，但也会带来精度波动。在移动端做意图分类这类任务时，精度下降通常可接受；若是做相似度匹配，建议保留16位量化。另外，GraphModel不支持动态形状，输入长度必须固定，这要求在预处理阶段就把序列补齐或截断到128。
内存与推理延迟优化
手机内存有限，TensorFlow.js默认会缓存中间张量，如果不手动释放，连续推理几次就会崩。MobileBERT在128长度下，16位量化模型权重约占30MB，但推理时产生的临时张量可能翻倍。正确做法是在每次executeAsync之后调用tf.dispose或把结果包在tf.tidy里。以下示例展示如何安全执行单句推理并回收资源。
async function infer(model, inputIds, attnMask) {
return tf.tidy(() => {
const ids = tf.tensor2d(inputIds, [1, 128], 'int32');
const mask = tf.tensor2d(attnMask, [1, 128], 'int32');
const out = model.executeAsync({ input_ids: ids, attention_mask: mask });
return out;
});
}
延迟方面，中端安卓机用tfjs-node单句推理约需400到700毫秒。如果走React Native的JavaScript桥，还要加上序列化开销。一种有效手段是把模型放在Web Worker或原生模块里，Node.js侧只做调度。另外，MobileBERT的层数为24，比BERT_base还深，因此减少序列长度是降延迟最直接的方式，业务允许时把128降到64能省近一半时间。
还有一个隐蔽问题是CPU线程数。TensorFlow.js底层使用Eigen，默认线程数可能等于核心数，在手机上引发发热降频。通过环境变量或原生配置限制为2到3线程，反而能让延迟更稳定。下表列出不同配置在红米低端机上的表现：
| 量化位宽 | 序列长度 | 线程数 | 平均延迟 |
|---|---|---|---|
| 32位 | 128 | 4 | 920ms |
| 16位 | 128 | 2 | 460ms |
| 16位 | 64 | 2 | 240ms |
中文分词与输入构造
MobileBERT预训练用的是WordPiece分词，中文需先按字切再查词表。Node.js没有官方中文BERT分词包，我们可以用bert-tokenization这类轻量库，或者自己读vocab.txt实现。核心是先把文本转成Unicode字符数组，再对每个字判断是否可继续拼接成词，最后加[CLS]和[SEP]。下面代码给出一个简化版分词函数。
function tokenize(text, vocab) {
const tokens = ['[CLS]'];
for (const ch of text) {
const piece = ch;
if (vocab[piece] !== undefined) {
tokens.push(piece);
} else {
tokens.push('[UNK]');
}
}
tokens.push('[SEP]');
const ids = tokens.map(t => vocab[t] !== undefined ? vocab[t] : vocab['[UNK]']);
return ids;
}
构造输入时还要生成attention_mask，全1即可，以及token_type_ids单句全0。注意MobileBERT某些版本不需要token_type_ids，传了反而报错，要对照模型签名。对于移动端，词表文件约80KB可放本地，避免每次网络拉取。若做实时语音转写后的语义理解，分词必须控制在10毫秒内，上述按字查表方式完全够用。
实际项目中，我们把分词和张量构造放到Node.js端，模型推理也留在同一进程，通过消息队列接收Native层发来的文本。这样避免了跨语言对象转换，整体在千元机上能稳定支撑每三秒一次的请求。如果未来要上多语言，只需替换vocab.txt与分词正则，模型结构无需变动。
Node.jsMobileBERTon-device_inference修改时间：2026-08-13 19:06:32