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

根据文档，该功能提供给 tf.tidy "must not return a 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>

Answer 2

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

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

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

Memory leak in Tensorflow.js: How to clean up unused tensors?

javascript

node.js

machine-learning

memory-leaks

tensorflow.js