Node.js环境下TensorFlow.js模型如何批量加载图片构建推理输入张量

实现核心逻辑 #

你现在单张图处理完成后得到的张量形状是(imageHeight, imageWidth, 3)，要生成符合模型输入要求的(batchNumber, imageHeight, imageWidth, 3)批量张量，核心操作就是把所有单图处理后的张量，在新增的第0维做统一堆叠。必须保证所有单图的高、宽、通道数完全一致，否则会出现形状不匹配的报错。

可落地的实现方案 #

方案1：逐张处理后堆叠（适配自定义预处理逻辑） #

适合需要对单张图做差异化处理的场景，完全兼容你已经写好的单图加载流程：

1. 遍历所有待加载的图片路径列表
2. 逐张执行你已有的读文件、解码、类型转换、自定义预处理流程
3. 对所有处理完成的单图张量统一做维度堆叠，生成最终批量输入张量

代码示例：

const tf = require('@tensorflow/tfjs-node');
const fs = require('fs/promises');
const path = require('path');

// 替换为你模型实际要求的输入尺寸
const MODEL_INPUT_WIDTH = 224;
const MODEL_INPUT_HEIGHT = 224;

async function loadBatchFromDir(imageDirPath) {
  // 读取目录下所有图片文件
  const allFiles = await fs.readdir(imageDirPath);
  const jpgImagePaths = allFiles
    .filter(fileName => /\.(jpe?g|png)$/i.test(fileName))
    .map(fileName => path.join(imageDirPath, fileName));

  const singleTensors = [];
  for (const imgPath of jpgImagePaths) {
    // 完全沿用你已有的单图加载逻辑
    const imgBuffer = await fs.readFile(imgPath);
    let imgTensor = tf.node.decodeImage(new Uint8Array(imgBuffer), 3);
    
    // 统一做尺寸对齐，确保所有图尺寸一致
    imgTensor = tf.image.resizeBilinear(imgTensor, [MODEL_INPUT_HEIGHT, MODEL_INPUT_WIDTH]);
    imgTensor = tf.cast(imgTensor, 'float32');

    // 这里加你原本写的其他图像处理逻辑，比如像素归一化、通道值调整等
    // imgTensor = imgTensor.div(255.0)

    singleTensors.push(imgTensor);
  }

  // 在第0维堆叠所有单图，最终形状自动变为 (batchCount, height, width, 3)
  const batchInputTensor = tf.stack(singleTensors, axis = 0);

  // 及时释放不再需要的单图张量，避免Node进程内存泄漏
  singleTensors.forEach(tensor => tensor.dispose());

  return batchInputTensor;
}

方案2：用自动内存管理的加载逻辑（大批量场景更高效） #

如果你的图片不需要特殊的差异化处理，可以结合tf.tidy作用域自动管理内存，省去手动释放张量的代码开销，适合处理上百张以上图片的批量任务：

async function fastLoadBatch(imagePathList) {
  // 并行读取所有图片文件，减少IO等待时间
  const imageBuffers = await Promise.all(
    imagePathList.map(imgPath => fs.readFile(imgPath))
  );

  // tf.tidy 会自动清理作用域内创建的所有非返回值张量
  const batchTensor = tf.tidy(() => {
    const processedTensors = imageBuffers.map(buffer => {
      const img = tf.node.decodeImage(new Uint8Array(buffer), 3);
      return tf.image.resizeBilinear(img, [MODEL_INPUT_HEIGHT, MODEL_INPUT_WIDTH]);
    });
    return tf.stack(processedTensors, axis = 0);
  });

  return batchTensor;
}

常见踩坑提醒 #

• 所有单图必须做统一尺寸缩放，否则不同分辨率的图片生成的张量形状不一致，tf.stack会直接抛出维度不匹配的错误
• 注意图像尺寸的维度顺序：TensorFlow.js的图像张量默认是[height, width, channel]格式，调用tf.image.resizeBilinear时传入的尺寸参数顺序是[height, width]，搞反顺序会直接导致推理精度大幅下降
• 如果你的原Keras模型训练时用的是channels_first的输入格式（即输入形状为(batch, 3, height, width)），堆叠完成后需要调用tf.transpose(batchTensor, [0, 3, 1, 2])把通道维度移动到正确位置，否则加载模型时会报输入形状不匹配的错误
• 不要一次性把过大数量的图片加载为单个张量，建议根据服务器内存大小控制单批图片数量（通常16/32张为一批），单批推理完成后及时释放对应的张量，避免长时间运行后内存溢出

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