Open
Description
System information
- Have I written custom code (as opposed to using a stock example script provided in TensorFlow.js):
- OS Platform and Distribution: Linux Ubuntu 16.04
- TensorFlow.js installed from (npm or script link): TFJS 4.4.0
Describe the current behavior
I was looking into creating TensorFlow models that shared weights between certain layers (called Siamese Networks), and I found that if a dense layer is applied to two different inputs in the same model, and the model is later disposed, the layer doesn't get disposed causing a memory leak.
Describe the expected behavior
The model should dispose layers that get used multiple times.
Standalone code to reproduce the issue
const tf = require("@tensorflow/tfjs");
console.log(`start: ${tf.memory().numTensors}`)
const input1 = tf.layers.input({shape: [1]});
const input2 = tf.layers.input({shape: [1]});
const denseLayer = tf.layers.dense({units: 1});
const output1 = denseLayer.apply(input1);
const output2 = denseLayer.apply(input2);
console.log(`layers: ${tf.memory().numTensors}`)
const model = tf.model({inputs: [input1, input2], outputs: [output1, output2]});
console.log(`model: ${tf.memory().numTensors}`)
model.dispose();
console.log(`dispose: ${tf.memory().numTensors}`)
denseLayer.dispose();
console.log(`dispose layer: ${tf.memory().numTensors}`)
Output:
start: 0
layers: 2
model: 2
dispose: 2
dispose layer: 0
Other info / logs
This can be easily resolved by combining the two inputs and only calling Dense.apply()
once, so I'm not sure if this is intended usage. A fix would be appreciated.