Tfjs: Model.dispose() misses some Tensors it seems

[rlexa]

TensorFlow.js version 0.14.1

Browser version Chrome 70.0.3538.77

Describe the problem or feature request

Expecting Model.dispose() to dispose all children, but compiling a model seems to produce at least a single Tensor which isn't disposed by the Model and can't be disposed manually.

Code to reproduce the bug / link to feature request

Case 1: create, dispose

This one seems to be working fine as in fact no Tensors are created anyway.

console.log('start', Environment.memory().numTensors); // start 0
const m = sequential();
console.log('+ model', Environment.memory().numTensors); // + model 0
console.log(m.dispose()); // {refCountAfterDispose: 0, numDisposedVariables: 0}
console.log('- model', Environment.memory().numTensors); // - model 0

Case 2: create, add layer, dispose

Adding a single Layer and disposing already leaks as it seems that the Layer isn't disposed properly in the Model.

console.log('start', Environment.memory().numTensors); // start 0
const m = sequential();
console.log('+ model', Environment.memory().numTensors); // + model 0
m.add(layers.dense({ units: 1, inputShape: [1] }));
console.log('+ layer', Environment.memory().numTensors); // + layer 4
console.log(m.dispose()); // {refCountAfterDispose: 0, numDisposedVariables: 2}
console.log('- model', Environment.memory().numTensors); // - model 2

Case 2: fixed with tidy() create, add layer, dispose

Wrapping the Layer addition in tidy() helps here, but as everything happens inside of Model.add() functions it feels like it shouldn't be necessary?

console.log('start', Environment.memory().numTensors); // start 0
const m = sequential();
console.log('+ model', Environment.memory().numTensors); // + model 0
tidy(() => {
  m.add(layers.dense({ units: 1, inputShape: [1] }));
});
console.log('+ layer', Environment.memory().numTensors); // + layer (tidy) 2
console.log(m.dispose()); // {refCountAfterDispose: 0, numDisposedVariables: 2}
console.log('- model', Environment.memory().numTensors); // - model 0

Case 3: create, add layer, compile, dispose

Model.compile() adds another not-disposed Tensor.

console.log('start', Environment.memory().numTensors); // start 0
const m = sequential();
console.log('+ model', Environment.memory().numTensors); // + model 0
tidy(() => {
  m.add(layers.dense({ units: 1, inputShape: [1] }));
});
console.log('+ layer', Environment.memory().numTensors); // + layer (tidy) 2
m.compile({ loss: 'meanSquaredError', optimizer: 'sgd' });
console.log('+ compile', Environment.memory().numTensors); // + compile 3
console.log(m.dispose()); // {refCountAfterDispose: 0, numDisposedVariables: 2}
console.log('- model', Environment.memory().numTensors); // - model 1

Case 3: failed to fix with tidy() create, add layer, compile, dispose

Wrapping compile() in a tidy() feels like it shouldn't be necessary in the first place, but anyway doesn't help.

console.log('start', Environment.memory().numTensors); // start 0
const m = sequential();
console.log('+ model', Environment.memory().numTensors); // + model 0
tidy(() => {
  m.add(layers.dense({ units: 1, inputShape: [1] }));
});
console.log('+ layer', Environment.memory().numTensors); // + layer (tidy) 2
tidy(() => {
  m.compile({ loss: 'meanSquaredError', optimizer: 'sgd' });
});
console.log('+ compile (tidy)', Environment.memory().numTensors); // + compile (tidy) 3
console.log(m.dispose()); // {refCountAfterDispose: 0, numDisposedVariables: 2}
console.log('- model', Environment.memory().numTensors); // - model 1

Case 4: create, add layer, compile, train, dispose

Training with Model.fit() works fine i.e. doesn't add more non-disposable Tensors if you don't forget to dispose the inputs manually (according to docs tidy() should not contain async code so can't wrap fit()).

console.log('start', Environment.memory().numTensors); // start 0
const m = sequential();
console.log('+ model', Environment.memory().numTensors); // + model 0
tidy(() => {
  m.add(layers.dense({ units: 1, inputShape: [1] }));
});
console.log('+ layer', Environment.memory().numTensors); // + layer (tidy) 2
tidy(() => {
  m.compile({ loss: 'meanSquaredError', optimizer: 'sgd' });
});
console.log('+ compile (tidy)', Environment.memory().numTensors); // + compile (tidy) 3

const xs = tensor2d([-1, 0, 1, 2, 3, 4], [6, 1]);
const ys = tensor2d([-3, -1, 1, 3, 5, 7], [6, 1]);
await m.fit(xs, ys, { epochs: 250 });
xs.dispose();
ys.dispose();
console.log('+ train', Environment.memory().numTensors); // + train 3

console.log(m.dispose()); // {refCountAfterDispose: 0, numDisposedVariables: 2}
console.log('- model', Environment.memory().numTensors); // - model 1

Case 5: create, add layer, compile, train, predict, dispose

Prediction with Model.predict() works fine i.e. doesn't add more non-disposable Tensors if you don't forget to dispose the output Tensor of the tidy() function.

console.log('start', Environment.memory().numTensors); // start 0
const m = sequential();
console.log('+ model', Environment.memory().numTensors); // + model 0
tidy(() => {
  m.add(layers.dense({ units: 1, inputShape: [1] }));
});
console.log('+ layer', Environment.memory().numTensors); // + layer (tidy) 2
tidy(() => {
  m.compile({ loss: 'meanSquaredError', optimizer: 'sgd' });
});
console.log('+ compile (tidy)', Environment.memory().numTensors); // + compile (tidy) 3

const xs = tensor2d([-1, 0, 1, 2, 3, 4], [6, 1]);
const ys = tensor2d([-3, -1, 1, 3, 5, 7], [6, 1]);
await m.fit(xs, ys, { epochs: 250 });
xs.dispose();
ys.dispose();
console.log('+ train', Environment.memory().numTensors); // + train 3

const prediction = tidy(() => m.predict(tensor2d([20], [1, 1]))) as Tensor;
console.log('...predicted', await prediction.data()); // ...predicted Float32Array [38.493194580078125]
prediction.dispose();
console.log('+ predict', Environment.memory().numTensors); // + predict 3

console.log(m.dispose()); // {refCountAfterDispose: 0, numDisposedVariables: 2}
console.log('- model', Environment.memory().numTensors); // - model 1

[justadudewhohacks]

Facing the same issue. I think the problem resides in the LayerVariable class. When initializing the tensor values of a layer, LayerVariables are created with the output of the initializer function. A new tfc.variable is created here, but the original tensor is not being disposed, causing a memory leak.

I think in some of the earlier releases you could simply create a variable from a tensor, but seems as if you now have to dispose the original tensor, or am I wrong with this assumption?

[piercus]

I'm facing the same issue.

I have tested this on 0.9.2

Case 3

console.log('start', Environment.memory().numTensors); // start 0
const m = sequential();
console.log('+ model', Environment.memory().numTensors); // + model 0
tidy(() => {
  m.add(layers.dense({ units: 1, inputShape: [1] }));
});
console.log('+ layer', Environment.memory().numTensors); // + layer (tidy) 2
m.compile({ loss: 'meanSquaredError', optimizer: 'sgd' });
console.log('+ compile', Environment.memory().numTensors); // + compile 3
console.log(m.dispose()); // {refCountAfterDispose: 0, numDisposedVariables: 2}
console.log('- model', Environment.memory().numTensors); // - model 1

is not fixed

Looking to the code, it seems that tensorflow/tfjs-layers#402 is part of 0.9.2

@caisq should we repoen this ?

Thank you for your help

EDIT Workaround is to do m.optimizer.dispose() then m.dispose()

[dsmilkov]

Thanks. We will fix this by making model.dispose() internally call model.optimizer.dispose() if the optimizer is not owned by the user. The optimizer is not owned by the user if the user passed a string (e.g. sgd) when calling model.compile() instead of an instance.