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()

Answer 2

根据文档，提供给tf.tidy的函数“不得返回 Promise”。 在内部，tf 后端处理拟合模型时使用的所有张量。 这就是为什么tf.fit不应该放在tf.tidy里面的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>