TensorFlow.js内存泄漏问题：如何清理未使用的Tensor？

我在编写TensorFlow.js训练脚本时遇到了内存泄漏问题：当神经网络训练崩溃（比如示例中故意设置无效的reshape层）时，训练被中断导致Tensor无法正常释放，进而引发内存泄漏。

目前tf.tidy仅支持同步函数调用，无法处理await model.fitDataset()这类异步场景。请问有没有办法清理未使用的Tensor？或者无需刷新页面就能释放所有现有Tensor的办法？

示例代码如下：

async function train(shouldCrash) {
 console.log(`Training, shouldCrash=${shouldCrash}`);
 const dataset = tf.data.zip({
 // setup data
 xs: tf.data.array([[1],[1]]),
 ys: tf.data.array([1]),
 }).batch(1);
 const model = tf.sequential({
 // setup model
 layers: [
 tf.layers.dense({units: 1, inputShape: [1]}),
 tf.layers.reshape({targetShape: [(shouldCrash ? 2 : 1)]}), // use invalid shape when crashing
 ],
 });
 model.compile({
 optimizer: 'sgd',
 loss: 'meanSquaredError'
 });
 console.log(' Tensors before:', tf.memory().numTensors);
 try {
 const history = await model.fitDataset(dataset, { epochs: 1 });
 } catch (err) {
 console.log(` Error: ${err.message}`);
 }
 console.log(' Tensors after:', tf.memory().numTensors);
}
(async () => {
 await train(false); // normal training
 await train(true); // training with error
})();

<script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@1.1.2/dist/tf.min.js"></script>

在这个示例中，第一次正常训练没有Tensor泄漏；但第二次强制训练崩溃后，数据集和模型相关的Tensor没有被正确释放。

好问题！TensorFlow.js的tf.tidy确实只覆盖同步执行的代码块，异步操作（比如fitDataset）中创建的Tensor不在其自动清理范围内，这也是你遇到崩溃后泄漏的核心原因。下面分享几个可行的解决方案：

1. 手动清理关键资源（推荐） #

在训练流程的finally块中手动清理模型、张量等资源，不管训练成功还是失败都能确保释放。这是最可控的方式，适合大多数场景。

修改你的示例代码，添加资源清理逻辑：

async function train(shouldCrash) {
 console.log(`Training, shouldCrash=${shouldCrash}`);
 let dataset = null;
 let model = null;
 console.log(' Tensors before:', tf.memory().numTensors);
 try {
  dataset = tf.data.zip({
   xs: tf.data.array([[1],[1]]),
   ys: tf.data.array([1]),
  }).batch(1);
  model = tf.sequential({
   layers: [
    tf.layers.dense({units: 1, inputShape: [1]}),
    tf.layers.reshape({targetShape: [(shouldCrash ? 2 : 1)]}),
   ],
  });
  model.compile({
   optimizer: 'sgd',
   loss: 'meanSquaredError'
  });
  const history = await model.fitDataset(dataset, { epochs: 1 });
 } catch (err) {
  console.log(` Error: ${err.message}`);
 } finally {
  // 清理模型：释放模型的所有层、变量和内部张量
  if (model) {
   model.dispose();
  }
  // 如果你手动创建了独立的Tensor（而非模型或数据集内部的），也要单独调用tf.dispose()
  console.log(' Tensors after cleanup:', tf.memory().numTensors);
 }
}
(async () => {
 await train(false); // normal training
 await train(true); // training with error
})();

2. 清理所有模型变量 #

如果你的场景中只需要清理模型相关的可训练变量，可以使用tf.disposeVariables()。这个方法会释放所有通过tf.variable()创建的张量，包括模型的权重。不过要注意：调用后所有现有模型都无法再使用，需要重新初始化。

示例用法：

// 在训练结束或崩溃后调用
tf.disposeVariables();

3. 全局重置后端（彻底清理） #

如果你需要一次性释放TensorFlow.js占用的所有内存（包括所有张量、模型、缓存等），可以使用tf.clearBackend()。这相当于完全重置TensorFlow.js的运行环境，适合需要彻底清理的场景（比如训练任务完成后、页面需要复用内存时）。

示例用法：

// 彻底清理所有资源
tf.clearBackend();
// 注意：之后需要重新初始化TensorFlow.js后端（通常会自动触发）
tf.ready().then(() => {
 // 可以重新创建模型和开始训练
});

额外提示 #

• 用tf.memory()实时监控张量数量，确认泄漏是否被解决。
• 对于异步流程中创建的临时张量，尽量在同步代码块中用tf.tidy包裹（比如数据预处理的同步部分），减少手动清理的负担。

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