TensorFlow.js内存泄漏：如何管理tf.data.generator创建的大型数据集内存？

问题根因 #

• tf.data.generator生成的张量默认不会自动释放，流水线消费后没有触发dispose逻辑
• 生成张量的过程中产生的大量中间张量没有被清理，累积占用内存
• 拆分图像、标签两个独立生成器会额外增加张量持有开销

修复方案 #

1. 清理生成张量时的中间变量 #

使用tf.tidy()包裹张量生成逻辑，除了需要返回/产出的最终张量外，所有中间创建的张量会被自动释放。以图像生成方法为例：

// Convert image buffer to a Tensor object
#generateTensor = (imagePath) => {
  return tf.tidy(() => {
    const buffer = fs.readFileSync(imagePath);
    return tf.node.decodeJpeg(buffer, 3)
        .resizeNearestNeighbor([128, 128])
        .toFloat()
        .div(tf.scalar(255.0))
  })
}

2. 合并图像、标签生成逻辑 #

避免拆分两个独立生成器重复遍历路径，减少不必要的张量持有开销：

// 合并后的生成器，直接产出样本+标签组合
* #dataGenerator(){
  for(let i=0; i<this.paths.length; ++i){
    let image;
    try {
      image = this.#generateTensor(this.paths[i]);
    } catch (error) {
      continue;
    }
    const classIndex = (path.basename(path.dirname(this.paths[i])) === 'Cat' ? 0 : 1);
    // 用tf.oneHot替代手动生成数组，减少内存开销
    const label = tf.oneHot(classIndex, 2, 'int32');
    yield {xs: image, ys: label};
  }
}

3. 优化数据集流水线配置 #

新增预加载逻辑提升性能，同时避免不必要的张量缓存：

// Load data
loadData = () => {
  console.log('\n\nLoading data...')
  const ds = tf.data.generator(this.#dataGenerator.bind(this))
      .batch(32)
      .shuffle(32)
      // 预加载1个batch，同时确保消费后的张量自动回收
      .prefetch(1);
  return ds;
}

4. 训练阶段主动释放batch #

如果是自定义迭代训练，每步训练完成后手动释放当前batch的所有张量：

for await (const batch of trainDS) {
  // 训练逻辑
  await model.trainOnBatch(batch.xs, batch.ys);
  // 释放当前batch的所有张量
  tf.dispose(batch);
}

如果使用内置的model.fitDataset()方法，TF.js会自动完成消费后张量的清理，无需额外手动操作。

验证方法 #

可通过tf.memory().numTensors打印当前内存中的张量数量，正常运行时数值会在固定区间波动，不会出现持续上涨的情况。

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