Added IDisposable support for several classes

src/Microsoft.ML.Data/DataLoadSave/CompositeDataLoader.cs

@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
+using System;
 using Microsoft.ML;
 using Microsoft.ML.Data;
 using Microsoft.ML.Runtime;
@@ -15,7 +16,7 @@ namespace Microsoft.ML.Data
     /// This class represents a data loader that applies a transformer chain after loading.
     /// It also has methods to save itself to a repository.
     /// </summary>
-    public sealed class CompositeDataLoader<TSource, TLastTransformer> : IDataLoader<TSource>
+    public sealed class CompositeDataLoader<TSource, TLastTransformer> : IDataLoader<TSource>, IDisposable
         where TLastTransformer : class, ITransformer
     {
         internal const string TransformerDirectory = TransformerChain.LoaderSignature;
@@ -110,5 +111,19 @@ private static VersionInfo GetVersionInfo()
                 loaderSignature: LoaderSignature,
                 loaderAssemblyName: typeof(CompositeDataLoader<,>).Assembly.FullName);
         }
+
+        #region IDisposable Support
+        private bool _disposed;
+
+        public void Dispose()
+        {
+            if (_disposed)
+                return;
+
+                Transformer.Dispose();
+
+            _disposed = true;
+        }
+        #endregion
     }
 }

src/Microsoft.ML.Data/DataLoadSave/TransformerChain.cs

@@ -50,7 +50,7 @@ internal interface ITransformerChainAccessor
     /// A chain of transformers (possibly empty) that end with a <typeparamref name="TLastTransformer"/>.
     /// For an empty chain, <typeparamref name="TLastTransformer"/> is always <see cref="ITransformer"/>.
     /// </summary>
-    public sealed class TransformerChain<TLastTransformer> : ITransformer, IEnumerable<ITransformer>, ITransformerChainAccessor
+    public sealed class TransformerChain<TLastTransformer> : ITransformer, IEnumerable<ITransformer>, ITransformerChainAccessor, IDisposable
     where TLastTransformer : class, ITransformer
     {
         private readonly ITransformer[] _transformers;
@@ -232,6 +232,21 @@ IRowToRowMapper ITransformer.GetRowToRowMapper(DataViewSchema inputSchema)
             }
             return new CompositeRowToRowMapper(inputSchema, mappers);
         }
+
+        #region IDisposable Support
+        private bool _disposed;
+
+        public void Dispose()
+        {
+            if (_disposed)
+                return;
+
+            foreach (var transformer in _transformers)
+                (transformer as IDisposable)?.Dispose();
+
+            _disposed = true;
+        }
+        #endregion
     }
 
     /// <summary>

src/Microsoft.ML.Data/DataView/CompositeRowToRowMapper.cs

@@ -14,7 +14,7 @@ namespace Microsoft.ML.Data
     /// A row-to-row mapper that is the result of a chained application of multiple mappers.
     /// </summary>
     [BestFriend]
-    internal sealed class CompositeRowToRowMapper : IRowToRowMapper
+    internal sealed class CompositeRowToRowMapper : IRowToRowMapper, IDisposable
     {
         [BestFriend]
         internal IRowToRowMapper[] InnerMappers { get; }
@@ -118,5 +118,20 @@ public SubsetActive(DataViewRow row, Func<int, bool> pred)
             /// </summary>
             public override bool IsColumnActive(DataViewSchema.Column column) => _pred(column.Index);
         }
+
+        #region IDisposable Support
+        private bool _disposed;
+
+        public void Dispose()
+        {
+            if (_disposed)
+                return;
+
+            foreach (var mapper in InnerMappers)
+                (mapper as IDisposable)?.Dispose();
+
+            _disposed = true;
+        }
+        #endregion
     }
 }

src/Microsoft.ML.Data/Prediction/PredictionEngine.cs

@@ -157,18 +157,13 @@ private protected virtual Func<DataViewSchema, IRowToRowMapper> TransformerCheck
         }
 
         public void Dispose()
-        {
-            Disposing(true);
-            GC.SuppressFinalize(this);
-        }
-
-        [BestFriend]
-        private protected void Disposing(bool disposing)
         {
             if (_disposed)
                 return;
-            if (disposing)
-                _disposer?.Invoke();
+
+            _disposer?.Invoke();
+            (Transformer as IDisposable)?.Dispose();
+
             _disposed = true;
         }
 

src/Microsoft.ML.Data/Scorers/MulticlassClassificationScorer.cs

@@ -64,7 +64,8 @@ private static VersionInfo GetVersionInfo()
         /// </summary>
         // REVIEW: It seems like the attachment of metadata should be solvable in a manner
         // less ridiculously verbose than this.
-        public sealed class LabelNameBindableMapper : ISchemaBindableMapper, ICanSaveModel, IBindableCanSavePfa, IBindableCanSaveOnnx
+        public sealed class LabelNameBindableMapper : ISchemaBindableMapper, ICanSaveModel, IBindableCanSavePfa,
+            IBindableCanSaveOnnx, IDisposable
         {
             private static readonly FuncInstanceMethodInfo1<LabelNameBindableMapper, object, Delegate> _decodeInitMethodInfo
                 = FuncInstanceMethodInfo1<LabelNameBindableMapper, object, Delegate>.Create(target => target.DecodeInit<int>);
@@ -379,6 +380,21 @@ public RowImpl(DataViewRow row, DataViewSchema schema)
                     public override ValueGetter<TValue> GetGetter<TValue>(DataViewSchema.Column column) => Input.GetGetter<TValue>(column);
                 }
             }
+
+            #region IDisposable Support
+            private bool _disposed;
+
+            public void Dispose()
+            {
+                // TODO: Is it necessary to call the base class Dispose()?
+                if (_disposed)
+                    return;
+
+                (_bindable as IDisposable)?.Dispose();
+
+                _disposed = true;
+            }
+            #endregion
         }
 
         /// <summary>

src/Microsoft.ML.Data/Scorers/PredictedLabelScorerBase.cs

@@ -17,7 +17,7 @@ namespace Microsoft.ML.Data
     /// Class for scorers that compute on additional "PredictedLabel" column from the score column.
     /// Currently, this scorer is used for binary classification, multi-class classification, and clustering.
     /// </summary>
-    internal abstract class PredictedLabelScorerBase : RowToRowScorerBase, ITransformCanSavePfa, ITransformCanSaveOnnx
+    internal abstract class PredictedLabelScorerBase : RowToRowScorerBase, ITransformCanSavePfa, ITransformCanSaveOnnx, IDisposable
     {
         public abstract class ThresholdArgumentsBase : ScorerArgumentsBase
         {
@@ -435,5 +435,20 @@ protected void EnsureCachedPosition<TScore>(ref long cachedPosition, ref TScore
                 cachedPosition = boundRow.Position;
             }
         }
+
+        #region IDisposable Support
+        private bool _disposed;
+
+        public void Dispose()
+        {
+            if (_disposed)
+                return;
+
+            (Bindings.RowMapper as IDisposable)?.Dispose();
+            (Bindable as IDisposable)?.Dispose();
+
+            _disposed = true;
+        }
+        #endregion
     }
 }

src/Microsoft.ML.Data/Scorers/PredictionTransformer.cs

@@ -40,7 +40,7 @@ internal static class PredictionTransformerBase
     /// Base class for transformers with no feature column, or more than one feature columns.
     /// </summary>
     /// <typeparam name="TModel">The type of the model parameters used by this prediction transformer.</typeparam>
-    public abstract class PredictionTransformerBase<TModel> : IPredictionTransformer<TModel>
+    public abstract class PredictionTransformerBase<TModel> : IPredictionTransformer<TModel>, IDisposable
         where TModel : class
     {
         /// <summary>
@@ -181,6 +181,22 @@ private protected void SaveModelCore(ModelSaveContext ctx)
                 }
             });
         }
+
+        #region IDisposable Support
+        private bool _disposed;
+
+        public void Dispose()
+        {
+            if (_disposed)
+                return;
+
+            (Model as IDisposable)?.Dispose();
+            (BindableMapper as IDisposable)?.Dispose();
+            (Scorer as IDisposable)?.Dispose();
+
+            _disposed = true;
+        }
+        #endregion
     }
 
     /// <summary>

src/Microsoft.ML.Data/Scorers/SchemaBindablePredictorWrapper.cs

@@ -32,7 +32,7 @@ namespace Microsoft.ML.Data
     /// This is a base class for wrapping <see cref="IPredictor"/>s in an <see cref="ISchemaBindableMapper"/>.
     /// </summary>
     internal abstract class SchemaBindablePredictorWrapperBase : ISchemaBindableMapper, ICanSaveModel, ICanSaveSummary,
-        IBindableCanSavePfa, IBindableCanSaveOnnx
+        IBindableCanSavePfa, IBindableCanSaveOnnx, IDisposable
     {
         // The ctor guarantees that Predictor is non-null. It also ensures that either
         // ValueMapper or FloatPredictor is non-null (or both). With these guarantees,
@@ -193,7 +193,7 @@ void ICanSaveSummary.SaveSummary(TextWriter writer, RoleMappedSchema schema)
         /// This class doesn't care. It DOES care that the role mapped schema specifies a unique Feature column.
         /// It also requires that the output schema has ColumnCount == 1.
         /// </summary>
-        protected sealed class SingleValueRowMapper : ISchemaBoundRowMapper
+        protected sealed class SingleValueRowMapper : ISchemaBoundRowMapper, IDisposable
         {
             private readonly SchemaBindablePredictorWrapperBase _parent;
 
@@ -241,7 +241,35 @@ DataViewRow ISchemaBoundRowMapper.GetRow(DataViewRow input, IEnumerable<DataView
                     getters[0] = _parent.GetPredictionGetter(input, InputRoleMappedSchema.Feature.Value.Index);
                 return new SimpleRow(OutputSchema, input, getters);
             }
+
+            #region IDisposable Support
+            private bool _disposed;
+
+            public void Dispose()
+            {
+                if (_disposed)
+                    return;
+
+                (_parent as IDisposable)?.Dispose();
+
+                _disposed = true;
+            }
+            #endregion
+        }
+
+        #region IDisposable Support
+        private bool _disposed;
+
+        public void Dispose()
+        {
+            if (_disposed)
+                return;
+
+            (Predictor as IDisposable)?.Dispose();
+
+            _disposed = true;
         }
+        #endregion
     }
 
     /// <summary>

src/Microsoft.ML.TensorFlow/TensorFlowModel.cs

@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
+using System;
 using Microsoft.ML.Runtime;
 using Microsoft.ML.TensorFlow;
 using Tensorflow;
@@ -13,7 +14,7 @@ namespace Microsoft.ML.Transforms
     /// It provides some convenient methods to query model schema as well as
     /// creation of <see cref="TensorFlowEstimator"/> object.
     /// </summary>
-    public sealed class TensorFlowModel
+    public sealed class TensorFlowModel : IDisposable
     {
         internal Session Session { get; }
         internal string ModelPath { get; }
@@ -31,6 +32,7 @@ internal TensorFlowModel(IHostEnvironment env, Session session, string modelLoca
             Session = session;
             ModelPath = modelLocation;
             _env = env;
+            _disposed = false;
         }
 
         /// <summary>
@@ -83,5 +85,19 @@ public TensorFlowEstimator ScoreTensorFlowModel(string outputColumnName, string
         /// </example>
         public TensorFlowEstimator ScoreTensorFlowModel(string[] outputColumnNames, string[] inputColumnNames, bool addBatchDimensionInput = false)
             => new TensorFlowEstimator(_env, outputColumnNames, inputColumnNames, this, addBatchDimensionInput);
+
+        #region IDisposable Support
+        private bool _disposed;
+
+        public void Dispose()
+        {
+            if (_disposed)
+                return;
+
+            Session.Dispose();
+
+            _disposed = true;
+        }
+        #endregion
     }
 }

src/Microsoft.ML.Vision/DnnRetrainTransform.cs

@@ -738,6 +738,8 @@ public void Dispose()
             {
                 if (_session != null && _session != IntPtr.Zero)
                 {
+                    if (_session.graph != null)
+                        _session.graph.Dispose();
                     _session.close();
                 }
             }

src/Microsoft.ML.Vision/ImageClassificationTrainer.cs

@@ -77,10 +77,8 @@ namespace Microsoft.ML.Vision
     /// </remarks>
     public sealed class ImageClassificationTrainer :
         TrainerEstimatorBase<MulticlassPredictionTransformer<ImageClassificationModelParameters>,
-            ImageClassificationModelParameters>, IDisposable
+            ImageClassificationModelParameters>
     {
-        private bool _isDisposed;
-
         internal const string LoadName = "ImageClassificationTrainer";
         internal const string UserName = "Image Classification Trainer";
         internal const string ShortName = "IMGCLSS";
@@ -1206,6 +1204,9 @@ private void UpdateTransferLearningModelOnDisk(int classCount)
             _session.graph.Dispose();
             _session.Dispose();
             _session = LoadTFSessionByModelFilePath(Host, frozenModelPath, false);
+
+            sess.graph.Dispose();
+            sess.Dispose();
         }
 
         private void VariableSummaries(RefVariable var)
@@ -1322,24 +1323,6 @@ private static TensorFlowSessionWrapper LoadTensorFlowSessionFromMetaGraph(IHost
             return new TensorFlowSessionWrapper(GetSession(env, modelFilePath, true), modelFilePath);
         }
 
-        public void Dispose()
-        {
-            if (_isDisposed)
-                return;
-
-            if (_session?.graph != IntPtr.Zero)
-            {
-                _session.graph.Dispose();
-            }
-
-            if (_session != null && _session != IntPtr.Zero)
-            {
-                _session.close();
-            }
-
-            _isDisposed = true;
-        }
-
         /// <summary>
         /// Trains a <see cref="ImageClassificationTrainer"/> using both training and validation data,
         /// returns a <see cref="ImageClassificationModelParameters"/>.

test/Microsoft.ML.AutoML.Tests/AutoFitTests.cs

@@ -54,8 +54,6 @@ public void AutoFitMultiTest()
         }
 
         [TensorFlowFact]
-        //Skipping test temporarily. This test will be re-enabled once the cause of failures has been determined
-        [Trait("Category", "SkipInCI")]
         public void AutoFitImageClassificationTrainTest()
         {
             var context = new MLContext(seed: 1);

test/Microsoft.ML.Core.Tests/UnitTests/TestEntryPoints.cs

@@ -6365,8 +6365,6 @@ public void TestOvaMacroWithUncalibratedLearner()
         }
 
         [TensorFlowFact]
-        //Skipping test temporarily. This test will be re-enabled once the cause of failures has been determined
-        [Trait("Category", "SkipInCI")]
         public void EntryPointTensorFlowTransform()
         {
             Env.ComponentCatalog.RegisterAssembly(typeof(TensorFlowTransformer).Assembly);
@@ -6382,8 +6380,6 @@ public void EntryPointTensorFlowTransform()
         }
 
         [TensorFlowFact]
-        //Skipping test temporarily. This test will be re-enabled once the cause of failures has been determined
-        [Trait("Category", "SkipInCI")]
         public void TestTensorFlowEntryPoint()
         {
             var dataPath = GetDataPath("Train-Tiny-28x28.txt");