Tensorflow.js中的内存泄漏：如何清理未使用的张量？

Question

我正在编写一段脚本，有时会泄漏张量。这可能发生在多种情况下，例如当我训练神经网络时，训练崩溃。在这种情况下，训练被中断，并且无法正确处理张量。这将导致内存泄漏，我正尝试通过处理未使用的张量来清理它。

示例

在下面的代码片段中，我正在训练两个（非常简单的）模型。第一次运行将正常工作，并且不会泄漏张量（训练前张量数=训练后张量数）。第二次，我使用一个无效的 reshape 层来强制训练过程崩溃。因此，会抛出错误并且来自数据集的张量将无法正确处理。这段代码是为了展示张量可能被泄漏而编写的示例。

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>

问题

有一个名为tf.tidy的函数，它可以帮助我在某些情况下处理未使用的张量，但它只能用于同步函数调用。因此，在调用await model.fitDataset(...)时无法使用。

是否有一种方法可以处理任何未使用的张量？或者，是否有一种方法可以处置页面上所有现有的张量（而不重新加载页面）？

- Thomas Dondorf

Answer 1

在异步代码中清除任何未使用的张量的方法是将创建它们的代码包在startScope()和endScope()调用之间。

tf.engine().startScope()
// do your thing
tf.engine().endScope()

- David

---

我通过查看 ts.tidy() 方法弄清楚了。TFjs 在他们的测试中也使用它 https://github.com/tensorflow/tfjs-core/blob/master/tfjs-core/src/jasmine_util.ts#L185。 - David

是的，太好了！它解决了我遇到的内存问题。谢谢！ - Juan

太棒了！！这真是神奇！它清理了开始和结束作用域之间的所有张量 :) - nk911

这个API有没有相关的文档？ - jameshfisher

TensorFlow.js的维护者建议使用.dispose()而不是上述方法：https://github.com/tensorflow/tfjs/issues/4685。如果您有多个一起运行的Promise，或者在Promise中创建Promise（这在async/await代码中很常见），那么上述方法可能会出现问题：这种情况可能会导致来自不同Promise的交错微任务，从而导致以下执行：startScope，startScope，endScope，endScope。第一个endScope可能会处理另一个Promise中将要使用的张量。 - Maxime Kjaer

Answer 2

根据文档，提供给 `tf.tidy` 的函数 "不能返回 Promise"。在拟合模型时，tf 后端会自动丢弃所有使用的张量。这就是为什么 `tf.fit` 不应该放在 `tf.tidy` 中。要销毁崩溃的模型，可以对模型调用 `tf.dispose`。

目前似乎存在内存泄漏问题，但是在定义模型期间出现模型崩溃是不良实现。这在适当的情况下不应发生，因为可以测试给定参数是否与层应输入的内容相匹配。例如，在构建模型之前可以避免将形状从 2 改变为 1，以防止内存泄漏。

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);
  return model
}

(async () => {
  const m1 = await train(false); // normal training
   tf.dispose(m1)
  const m2 = await train(true); // training with error
  
  tf.dispose(m2)
  tf.disposeVariables() 
  console.log('Tensors afters:', tf.memory().numTensors);
   
})();

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

- edkeveked

---

我知道我展示的代码是“糟糕的实现”，但正如我所说，它只是用来展示内存泄漏的。tf.disposeVariables函数看起来非常有用，我甚至不知道我可以将模型传递给tf.dispose。谢谢！ :) - Thomas Dondorf