TensorFlow JS报错：输入张量维度不匹配问题求助

问题场景 #

将已保存模型转换为TensorFlow JS格式后，用NodeJS加载model.json并创建图片预测接口，代码如下：

const tf = require("@tensorflow/tfjs-node");
var model;
loadModel();

async function loadModel() {
    model = await tf.loadGraphModel('https://LinkTo/YourModel/model.json');
    console.log("Model Loading Done!!")
}

async function detectcellPhone(imgurl) {
    const imgTensor = tf.node.decodeImage(new Uint8Array(fs.readFileSync(imgurl)), 3);
    const predictions = await model.executeAsync(imgTensor.expandDims(0));
    return predictions;
}

模型加载成功，但预测时触发错误：

Error: Invalid TF_Status: 3 Message: In[0] and In[1] has different ndims: [1,8,8,64,2] vs. [2,1]

错误原因 #

这个错误说明模型需要多个输入张量，但你只传入了一个，或者传入的输入张量形状和模型预期不匹配。报错里的两个维度[1,8,8,64,2]和[2,1]对应模型的两个输入，当前只传了第一个，导致模型找不到第二个输入，或者传入的第二个输入形状不符合要求。

解决步骤 #

1. 确认模型的输入要求
   在模型加载完成后，打印所有输入的详细信息，明确模型需要几个输入、每个输入的形状和类型：

   async function loadModel() {
       model = await tf.loadGraphModel('https://LinkTo/YourModel/model.json');
       console.log("Model Loading Done!!");
       // 遍历打印所有输入信息
       model.inputs.forEach((input, index) => {
           console.log(`输入${index+1}: 名称=${input.name}, 形状=${input.shape}, 类型=${input.dtype}`);
       });
   }
   

2. 修正图片预处理流程
   原代码的图片预处理缺少尺寸调整和归一化，这会导致输入形状和模型训练时不一致。需要和训练阶段的预处理逻辑对齐：

   • 调整图片尺寸到模型要求的输入大小（比如224x224、416x416，根据训练时的参数修改）
   • 对像素值做归一化（比如除以255，或者使用训练时的均值/标准差）

3. 补充缺失的输入张量
   如果模型确实需要多个输入，根据打印的输入信息构造对应的张量，一起传入executeAsync。比如如果第二个输入是形状[2,1]的张量，需要构造符合要求的张量传入。

修正后的完整代码示例 #

const tf = require("@tensorflow/tfjs-node");
const fs = require('fs'); // 原代码遗漏了fs模块导入
var model;

// 加载模型并打印输入信息
async function loadModel() {
    model = await tf.loadGraphModel('https://LinkTo/YourModel/model.json');
    console.log("Model Loading Done!!");
    // 打印模型输入要求
    model.inputs.forEach((input, index) => {
        console.log(`输入${index+1}: 名称=${input.name}, 形状=${input.shape}, 类型=${input.dtype}`);
    });
}

// 图片预测函数
async function detectcellPhone(imgurl) {
    // 读取图片文件
    const imgBuffer = fs.readFileSync(imgurl);
    // 解码图片为3通道张量
    let imgTensor = tf.node.decodeImage(imgBuffer, 3);
    
    // 1. 调整图片尺寸到模型要求的输入大小（示例为224x224，根据实际情况修改）
    imgTensor = tf.image.resizeBilinear(imgTensor, [224, 224]);
    // 2. 归一化（和训练时一致，示例为除以255）
    imgTensor = imgTensor.div(255.0);
    // 3. 增加batch维度（模型输入需要[batch, height, width, channels]）
    const input1 = imgTensor.expandDims(0);

    // 如果模型需要第二个输入，构造对应的张量（示例为形状[2,1]的张量，根据实际要求修改）
    const input2 = tf.tensor2d([[1], [2]], [2, 1]);

    // 传入所有输入张量执行预测
    const predictions = await model.executeAsync([input1, input2]);

    // 将张量转为普通数组返回，并释放张量内存避免泄漏
    const result = await predictions.data();
    // 释放所有临时张量
    imgTensor.dispose();
    input1.dispose();
    input2.dispose();
    predictions.dispose();

    return result;
}

// 启动加载模型
loadModel();

关键注意事项 #

• 必须保证输入张量的形状、数据类型、预处理逻辑和模型训练阶段完全一致
• 如果是目标检测类模型（如YOLO、SSD），通常需要额外输入锚框参数或先验框，需要对照训练代码确认参数值
• 记得释放张量内存，避免NodeJS内存泄漏

内容的提问来源于stack exchange，提问作者S M