Tensorflow.js(Node)训练MNIST数据集时频繁打印"Cleanup called"日志求助

核心问题：Tensor内存泄漏触发频繁清理 #

你的生成器每次yield出的imageTensor和oneHotTensor没有被手动释放，TensorFlow.js的自动内存回收机制会不断触发清理操作，这就是控制台持续打印"Cleanup called"的原因。此外代码还有几个可优化的点：

1. 手动释放Tensor内存 #

在yield后调用tensor.dispose()释放不再使用的Tensor，或用tf.tidy包裹Tensor创建逻辑，自动清理中间Tensor：

function* dataGenerator(type) {
  const dataRoot = `MNIST/${type}-Data`;
  const labels = fs.readdirSync(dataRoot);
  // 将标签转为数字，避免字符串作为数组索引的问题
  const labelNumbers = labels.map(l => parseInt(l));

  for (const label of labelNumbers) {
    const files = fs.readdirSync(`${dataRoot}/${label}`);
    for (const file of files) {
      const img = fs.readFileSync(`${dataRoot}/${label}/${file}`);
      // 用tf.tidy自动清理内部未保留的Tensor
      const result = tf.tidy(() => {
        // 归一化图像数据，适配神经网络训练
        const imageTensor = tf.node.decodePng(img, 1).cast('float32').div(255);
        const oneHotArr = new Array(labels.length).fill(0);
        oneHotArr[label] = 1;
        const oneHotTensor = tf.tensor1d(oneHotArr, 'float32');
        // 克隆要yield的Tensor，避免被tidy清理
        return { xs: imageTensor.clone(), ys: oneHotTensor.clone() };
      });
      yield result;
    }
  }
}

2. 改用异步文件读取提升性能 #

同步读取文件会阻塞事件循环，影响训练效率，改成异步读取配合异步生成器：

async function* dataGenerator(type) {
  const dataRoot = `MNIST/${type}-Data`;
  const labels = await fs.promises.readdir(dataRoot);
  const labelNumbers = labels.map(l => parseInt(l));

  for (const label of labelNumbers) {
    const files = await fs.promises.readdir(`${dataRoot}/${label}`);
    for (const file of files) {
      const img = await fs.promises.readFile(`${dataRoot}/${label}/${file}`);
      const result = tf.tidy(() => {
        const imageTensor = tf.node.decodePng(img, 1).cast('float32').div(255);
        const oneHotArr = new Array(labels.length).fill(0);
        oneHotArr[label] = 1;
        const oneHotTensor = tf.tensor1d(oneHotArr, 'float32');
        return { xs: imageTensor.clone(), ys: oneHotTensor.clone() };
      });
      yield result;
    }
  }
}

TensorFlow.js的tf.data.generator支持异步生成器，不会阻塞主线程。

3. 优化数据集配置 #

• 增大shuffle缓冲区：当前设置的100太小，针对60000条训练数据，建议设为1000或更大，提升数据打乱效果：

  const trainingDataset = tf.data.generator(() => dataGenerator("Training"))
    .shuffle(1000)
    .batch(100);
  

• 验证集无需shuffle：验证集顺序不影响评估结果，去掉该操作减少不必要的计算：

  const validationDataset = tf.data.generator(() => dataGenerator("Validation"))
    .batch(100);
  

4. 额外建议：归一化图像数据 #

训练前将图像张量归一化到[0,1]区间（.div(255)），是神经网络训练的常规操作，能加快模型收敛速度。

总结 #

修改后，Tensor内存会被及时释放，"Cleanup called"的日志会大幅减少，同时训练效率和模型收敛效果也会提升。

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