Skip to content

Commit a072b5e

Browse files
KsenijaScodemzs
authored andcommitted
Fix memory leak in TensorflowTransform (dotnet#4223)
* Fix memory leak by freeing input tensors after executing the graph * Fix memory leak by freeing the graph after each session * Get input tensors from runner
1 parent 7fd670f commit a072b5e

File tree

2 files changed

+17
-1
lines changed

2 files changed

+17
-1
lines changed

src/Microsoft.ML.Dnn/DnnUtils.cs

+5
Original file line numberDiff line numberDiff line change
@@ -399,6 +399,11 @@ public Runner AddInput(Tensor value, int index)
399399
return this;
400400
}
401401

402+
public List<Tensor> GetInputValues()
403+
{
404+
return _inputValues;
405+
}
406+
402407
public Runner AddOutputs(string output)
403408
{
404409
_outputs.Add(ParseOutput(output));

src/Microsoft.ML.TensorFlow/TensorflowTransform.cs

+12-1
Original file line numberDiff line numberDiff line change
@@ -463,6 +463,11 @@ private void Dispose(bool disposing)
463463
{
464464
Session.close(); // invoked Dispose()
465465
}
466+
467+
if (Session != null && Session.graph != IntPtr.Zero)
468+
{
469+
Session.graph.Dispose();
470+
}
466471
}
467472
finally
468473
{
@@ -645,7 +650,7 @@ private void UpdateCacheIfNeeded(long position, ITensorValueGetter[] srcTensorGe
645650
Runner runner = new Runner(_parent.Session);
646651

647652
// Feed inputs to the graph.
648-
for (int i = 0; i < _parent.Inputs.Length; i++)
653+
for (int i = 0; i < _parent.Inputs.Length; i++)
649654
{
650655
var tensor = srcTensorGetters[i].GetTensor();
651656
runner.AddInput(_parent.Inputs[i], tensor);
@@ -658,6 +663,12 @@ private void UpdateCacheIfNeeded(long position, ITensorValueGetter[] srcTensorGe
658663
// Execute the graph.
659664
var tensors = runner.Run();
660665

666+
List<Tensor> inputTensors = runner.GetInputValues();
667+
foreach (Tensor inputTensor in inputTensors)
668+
{
669+
inputTensor.Dispose();
670+
}
671+
661672
Contracts.Assert(tensors.Length > 0);
662673

663674
for (int j = 0; j < activeOutputColNames.Length; j++)

0 commit comments

Comments
 (0)