如何在TensorFlow.js中加载ONNX模型 仅用JS实现转换与推理

直接转换为TF.js原生格式实现 #

• 依赖纯JS实现的onnx-proto库解析ONNX文件结构，配合标准CNN算子映射逻辑即可完成转换，针对ResNet18/34、MobileNet这类常规图像分类模型，算子覆盖度完全满足需求：
  1. 接收ONNX二进制文件后，调用onnx-proto的decode方法解析出完整的模型计算图、权重参数、输入输出维度配置
  2. 执行算子映射：上述三类模型用到的Conv2d、BatchNormalization、ReLU、MaxPool、GlobalAveragePool、Dense、Softmax等算子，TF.js均有对应原生实现，直接按ONNX算子参数映射到TF.js层参数即可，例如ONNX的Conv2d的pads参数对应TF.js Conv2d层的padding配置，strides参数可直接复用
  3. 组装TF.js模型：按计算图的节点顺序拼接所有TF.js层，将解析得到的ONNX权重赋值到对应层的kernel、bias等参数位置，最终可导出为TF.js标准的model.json+分片权重二进制文件，也可直接在内存中生成可调用的TF.js模型实例
• 转换完成后直接用tf.loadLayersModel()或tf.loadGraphModel()加载即可执行推理

运行时兼容方案（无需持久化转换） #

如果不需要导出存储TF.js格式的模型文件，仅需要在JS环境中用TF.js API执行推理，可采用运行时适配方案：

• 用纯JS ONNX解析库读取计算图和权重，运行时直接调用TF.js的对应算子执行前向计算，无需导出为标准TF.js模型文件
• 图像分类场景的输入预处理、输出后处理逻辑可直接复用原TF.js业务代码，无需额外适配

注意事项 #

• 权重格式兼容：ONNX权重默认多为float32格式，和TF.js默认张量格式完全一致，直接转换为tf.Tensor即可使用，若遇到int8量化模型，先执行反量化操作再赋值
• 维度顺序对齐：ONNX的图像输入默认维度顺序为NCHW（batch、channel、height、width），TF.js默认采用NHWC顺序，转换时可调整输入张量维度，或在Conv2d层配置dataFormat: 'channelsFirst'直接兼容NCHW输入
• 公开预训练模型适配：如果使用公开预训练的ResNet、MobileNet ONNX模型，已有成熟的纯JS转换脚本可直接复用，无需从零实现算子映射逻辑

内容的提问来源于stack exchange，提问作者Matheus Campos