Tensorflow.js 中的内存泄漏:如何清理未使用的张量? [英] Memory leak in Tensorflow.js: How to clean up unused tensors?

问题描述

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

I'm writing a script, which sometimes leaks tensors. This can happen in multiple cases, for example when I'm training a neural network, but the training crashes. In this case, the training is interrupted and will not correctly dispose the tensors. This results in a memory leak, which I'm trying to clean up by disposing unused tensors.

示例

在下面的片段中，我正在训练两个(非常简单的)模型.第一次运行将起作用并且不会导致张量泄漏(训练前的张量数 = 训练后的张量数).第二次，我使用无效的 reshape 层在训练期间强制崩溃.因此，会抛出错误，并且数据集中的张量(我猜?)将不会被正确处理.该代码是展示张量可能如何泄漏的示例.

In the snippet below, I'm training two (very simple) models. The first run will work and will result in no leaked tensors (number of tensors before training = number of tensors after training). The second time, I'm using an invalid reshape layer to force a crash during training. Therefore, an error is thrown and the tensors from the dataset (I guess?) will not be correctly disposed. The code is an example to show how tensors might be leaked.

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(...)时不能使用.

有没有办法处理任何未使用的张量?或者，有没有办法处理页面上所有现有的张量(无需重新加载)?

Is there a way to dispose any unused tensors? Alternatively, is there a way to dispose all existing tensors on the page (without reloading it)?

推荐答案

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

The way to clean any unused tensors in async code is to wrap the code that creates them between a startScope() and an endScope() call.

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

这篇关于Tensorflow.js 中的内存泄漏:如何清理未使用的张量?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！