Skip to content

Calling Dense.apply() multiple times causes a memory leak #7667

Open
@taylorhansen

Description

@taylorhansen

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.

Metadata

Metadata

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions