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(...).有没有办法处理任何未使用的张量？或者，有没有办法处理页面上所有现有的张量（无需重新加载）？

Answer 1

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

tf.engine().startScope()

// do your thing

tf.engine().endScope()