TensorFlow.js加载自定义Keras模型时输入维度不匹配问题求助

嘿，这个问题我之前折腾TensorFlow.js模型的时候也踩过坑！其实核心原因很简单：你在Keras里训练模型的时候，默认是批量输入数据的，所以模型的输入层期望的是一个4维张量——[批次大小, 图像高度, 图像宽度, 通道数]，但tf.browser.fromPixels()从DOM元素拿到的是3维张量[255,255,3]，少了最前面的批次维度，这就导致了维度不匹配的报错。

快速解决方法 #

只需要给输入张量添加一个批次维度就行，用expandDims()方法最直观：

// 原来的代码
// const t = tf.browser.fromPixels(imgEl)

// 修改后：添加批次维度（在第0位插入维度，变成[1,255,255,3]）
const t = tf.browser.fromPixels(imgEl).expandDims(0);

这样处理后，输入张量的维度就和模型期望的一致了，预测就能正常运行。

完整修正后的代码示例 #

顺便给你补全一下代码，还要注意两个容易忽略的点：训练时的输入预处理要同步，以及清理张量避免内存泄漏：

async function runObjectDetection() {
  // 加载模型
  const net = await tf.loadLayersModel('http://localhost:8000/converted/model.json');
  
  // 获取页面上的图片元素
  const imgEl = document.getElementById('img');
  
  // 1. 转换为TensorFlow张量
  // 2. 添加批次维度（模型要求4维输入）
  // 3. 同步训练时的预处理：如果训练时把像素值除以255归一化，这里也要做！
  const inputTensor = tf.browser.fromPixels(imgEl)
    .expandDims(0)
    .div(255.0);
  
  // 执行预测
  let result = await net.predict(inputTensor);
  
  // 预测结果也是带批次维度的，用squeeze()去掉方便后续处理
  result = result.squeeze();
  
  // 这里可以添加结果处理逻辑，比如打印到控制台或者可视化
  console.log('检测结果数据:', await result.array());
  
  // 手动清理张量，避免浏览器内存泄漏
  inputTensor.dispose();
  result.dispose();
}

// 调用检测函数
runObjectDetection();

为什么你之前重塑没成功？ #

如果你之前尝试用reshape()，可能是写错了维度参数——应该写成reshape([1,255,255,3])而不是只保留原来的3维。不过expandDims(0)比手动reshape更直观，也不容易写错维度顺序。

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