我正在编写一段脚本,有时会泄漏张量。这可能发生在多种情况下,例如当我训练神经网络时,训练崩溃。在这种情况下,训练被中断,并且无法正确处理张量。这将导致内存泄漏,我正尝试通过处理未使用的张量来清理它。
示例
在下面的代码片段中,我正在训练两个(非常简单的)模型。第一次运行将正常工作,并且不会泄漏张量(训练前张量数=训练后张量数)。第二次,我使用一个无效的 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(...)
时无法使用。
是否有一种方法可以处理任何未使用的张量?或者,是否有一种方法可以处置页面上所有现有的张量(而不重新加载页面)?
.dispose()
而不是上述方法:https://github.com/tensorflow/tfjs/issues/4685。如果您有多个一起运行的Promise,或者在Promise中创建Promise(这在async/await代码中很常见),那么上述方法可能会出现问题:这种情况可能会导致来自不同Promise的交错微任务,从而导致以下执行:startScope
,startScope
,endScope
,endScope
。第一个endScope
可能会处理另一个Promise中将要使用的张量。 - Maxime Kjaer