Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fixes ort memory leak #5517

Closed
wants to merge 1 commit into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 25 additions & 3 deletions src/Microsoft.ML.OnnxTransformer/OnnxTransform.cs
Original file line number Diff line number Diff line change
Expand Up @@ -507,14 +507,32 @@ protected override Delegate MakeGetter(DataViewRow input, int iinfo, Func<int, b
}
}

private class OnnxRuntimeOutputCacher
private class OnnxRuntimeOutputCacher : IDisposable
{
public long Position;
public Dictionary<string, NamedOnnxValue> Outputs;
public Dictionary<string, DisposableNamedOnnxValue> Outputs;
public OnnxRuntimeOutputCacher()
{
Position = -1;
Outputs = new Dictionary<string, NamedOnnxValue>();
Outputs = new Dictionary<string, DisposableNamedOnnxValue>();
}

private bool _isDisposed;
public void Dispose()
{
if (_isDisposed)
return;
foreach (var onnxValue in Outputs.Values)
onnxValue.Dispose();
_isDisposed = true;
}

~OnnxRuntimeOutputCacher()
{
if (_isDisposed)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does it ever hit this code path?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The finalizer is hit. I dont think the Dispose call is though. I left it in cause it would be better to switch to a dispose call when we are able to.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We shouldn't be using a finalizer here.

  1. We should ensure OnnxRuntimeOutputCacher instances get disposed properly.
  2. In the cases where Dispose() is never called by the user, and this object is no longer referenced, all the references to the underlying DisposableNamedOnnxValue will be removed as well. If DisposableNamedOnnxValue is holding on to native memory, it should have a finalizer (or more preferrably use a SafeHandle). The managed object that manages the native memory is responsible for ensuring proper clean up.

We shouldn't be using finalizers to dispose managed objects. Finalizers should only be used for managing native resources.

return;
foreach (var onnxValue in Outputs.Values)
onnxValue.Dispose();
}
}

Expand All @@ -534,6 +552,10 @@ private void UpdateCacheIfNeeded(long position, INamedOnnxValueGetter[] srcNamed

foreach (var outputNameOnnxValue in outputNamedOnnxValues)
{
if(outputCache.Outputs.TryGetValue(outputNameOnnxValue.Name, out DisposableNamedOnnxValue value))
{
value.Dispose();
}
outputCache.Outputs[outputNameOnnxValue.Name] = outputNameOnnxValue;
}
outputCache.Position = position;
Expand Down
2 changes: 1 addition & 1 deletion src/Microsoft.ML.OnnxTransformer/OnnxUtils.cs
Original file line number Diff line number Diff line change
Expand Up @@ -350,7 +350,7 @@ public static OnnxModel CreateFromBytes(byte[] modelBytes, int? gpuDeviceId = nu
/// </summary>
/// <param name="inputNamedOnnxValues">The NamedOnnxValues to score.</param>
/// <returns>Resulting output NamedOnnxValues list.</returns>
public IReadOnlyCollection<NamedOnnxValue> Run(List<NamedOnnxValue> inputNamedOnnxValues)
public IDisposableReadOnlyCollection<DisposableNamedOnnxValue> Run(List<NamedOnnxValue> inputNamedOnnxValues)
{
return _session.Run(inputNamedOnnxValues);
}
Expand Down
12 changes: 6 additions & 6 deletions test/Microsoft.ML.Tests/OnnxConversionTest.cs
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,7 @@ public void SimpleEndToEndOnnxConversionTest()

// Step 3: Check ONNX model's text format. This test will be not necessary if Step 2 can run on Linux and
// Mac to support cross-platform tests.

CheckEquality(subDir, onnxTextName, digitsOfPrecision: 3);

Done();
Expand Down Expand Up @@ -139,7 +139,7 @@ private class BreastCancerBinaryClassification
[Fact]
public void KmeansOnnxConversionTest()
{
// Create a new context for ML.NET operations. It can be used for exception tracking and logging,
// Create a new context for ML.NET operations. It can be used for exception tracking and logging,
// as a catalog of available operations and as the source of randomness.
var mlContext = new MLContext(seed: 1);

Expand Down Expand Up @@ -384,7 +384,7 @@ public void TextNormalizingOnnxConversionTest()
new TextNormalizingEstimator(mlContext, keepDiacritics: true, caseMode: TextNormalizingEstimator.CaseMode.Upper, columns: new[] { ("UpperText", "text") })).Append(
new TextNormalizingEstimator(mlContext, keepDiacritics: true, caseMode: TextNormalizingEstimator.CaseMode.None, columns: new[] { ("OriginalText", "text") }));
var onnxFileName = $"TextNormalizing.onnx";

TestPipeline(pipeline, dataView, onnxFileName, new ColumnComparison[] { new ColumnComparison("NormText"), new ColumnComparison("UpperText"), new ColumnComparison("OriginalText") });

Done();
Expand Down Expand Up @@ -1154,7 +1154,7 @@ public void IndicateMissingValuesOnnxConversionTest()

// IsNaN outputs a binary tensor. Support for this has been added in the latest version
// of Onnxruntime, but that hasn't been released yet.
// So we need to convert its type to Int32 until then.
// So we need to convert its type to Int32 until then.
// ConvertType part of the pipeline can be removed once we pick up a new release of the Onnx runtime

var pipeline = mlContext.Transforms.IndicateMissingValues(new[] { new InputOutputColumnPair("MissingIndicator", "Features"), })
Expand Down Expand Up @@ -1806,7 +1806,7 @@ public void NonDefaultColNamesMultiClassificationOnnxConversionTest()
}
Done();
}

[Fact]
public void OneHotHashEncodingOnnxConversionWithCustomOpSetVersionTest()
{
Expand Down Expand Up @@ -2029,7 +2029,7 @@ private void TestPipeline<TLastTransformer, TRow>(EstimatorChain<TLastTransforme
private void TestPipeline<TLastTransformer>(EstimatorChain<TLastTransformer> pipeline, IDataView dataView, string onnxFileName, ColumnComparison[] columnsToCompare, string onnxTxtName = null, string onnxTxtSubDir = null)
where TLastTransformer : class, ITransformer
{
var model = pipeline.Fit(dataView);
using var model = pipeline.Fit(dataView);
var transformedData = model.Transform(dataView);
var onnxModel = ML.Model.ConvertToOnnxProtobuf(model, dataView);

Expand Down