Memory leak in Tensorflow.js: How to clean up unused tensors?
As per the documentation, the function provided to tf.tidy
"must not return a Promise". Internally, tf backend disposes all the tensors uses when fitting a model. That is why tf.fit
should not be placed inside tf.tidy
. To dispose the model crashed, one can call tf.dispose
on the model.
It is true that there seems to be currently a memory leak, but a model crash during the definition of the model is a poor implementation. This should not happen in a proper scenario as one can test if the parameters given matches what should be the input to the layers. For instance reshaping a shape of 2 to 1 can be avoid before constructing the model to prevent the memory leak.
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/[email protected]/dist/tf.min.js"></script>
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()