Add API option to store models on disk (instead of in memory); fix IEstimator memory leak (dotnet#269)

Author: daholste

src/Microsoft.ML.Auto/API/ExperimentSettings.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.IO;
 using System.Threading;
 
 namespace Microsoft.ML.Auto
@@ -12,6 +13,13 @@ public class ExperimentSettings
         public CancellationToken CancellationToken { get; set; } = default;
 
         /// <summary>
+        /// This is a pointer to a directory where all models trained during the AutoML experiment will be saved.
+        /// If null, models will be kept in memory instead of written to disk.
+        /// (Please note: for an experiment with high runtime operating on a large dataset, opting to keep models in 
+        /// memory could cause a system to run out of memory.)
+        /// </summary>
+        public DirectoryInfo ModelDirectory { get; set; } = null;
+
         /// This setting controls whether or not an AutoML experiment will make use of ML.NET-provided caching.
         /// If set to true, caching will be forced on for all pipelines. If set to false, caching will be forced off.
         /// If set to null (default value), AutoML will decide whether to enable caching for each model.

src/Microsoft.ML.Auto/API/RunResult.cs

@@ -3,14 +3,16 @@
 // See the LICENSE file in the project root for more information.
 
 using System;
+using System.IO;
 using System.Linq;
+using Microsoft.ML.Data;
 
 namespace Microsoft.ML.Auto
 {
     public sealed class RunResult<T>
     {
         public T ValidationMetrics { get; private set; }
-        public ITransformer Model { get; private set; }
+        public ITransformer Model { get { return _modelContainer.GetModel(); } }
         public Exception Exception { get; private set; }
         public string TrainerName { get; private set; }
         public int RuntimeInSeconds { get; private set; }
@@ -19,16 +21,17 @@ public sealed class RunResult<T>
         internal Pipeline Pipeline { get; private set; }
         internal int PipelineInferenceTimeInSeconds { get; private set; }
 
-        internal RunResult(
-            ITransformer model,
+        private readonly ModelContainer _modelContainer;
+
+        internal RunResult(ModelContainer modelContainer,
             T metrics,
             IEstimator<ITransformer> estimator,
             Pipeline pipeline,
             Exception exception,
             int runtimeInSeconds,
             int pipelineInferenceTimeInSeconds)
         {
-            Model = model;
+            _modelContainer = modelContainer;
             ValidationMetrics = metrics;
             Pipeline = pipeline;
             Estimator = estimator;

src/Microsoft.ML.Auto/Experiment/Experiment.cs

@@ -5,6 +5,8 @@
 using System;
 using System.Collections.Generic;
 using System.Diagnostics;
+using System.IO;
+using System.Linq;
 using System.Text;
 using Microsoft.Data.DataView;
 
@@ -22,9 +24,11 @@ internal class Experiment<T> where T : class
         private readonly ExperimentSettings _experimentSettings;
         private readonly IMetricsAgent<T> _metricsAgent;
         private readonly IEnumerable<TrainerName> _trainerWhitelist;
+        private readonly DirectoryInfo _modelDirectory;
 
         private IDataView _trainData;
         private IDataView _validationData;
+        private ITransformer _preprocessorTransform;
 
         List<RunResult<T>> iterationResults = new List<RunResult<T>>();
 
@@ -57,17 +61,17 @@ public Experiment(MLContext context,
             _experimentSettings = experimentSettings;
             _metricsAgent = metricsAgent;
             _trainerWhitelist = trainerWhitelist;
+            _modelDirectory = GetModelDirectory(_experimentSettings.ModelDirectory);
         }
 
         public List<RunResult<T>> Execute()
         {
-            ITransformer preprocessorTransform = null;
             if (_preFeaturizers != null)
             {
                 // preprocess train and validation data
-                preprocessorTransform = _preFeaturizers.Fit(_trainData);
-                _trainData = preprocessorTransform.Transform(_trainData);
-                _validationData = preprocessorTransform.Transform(_validationData);
+                _preprocessorTransform = _preFeaturizers.Fit(_trainData);
+                _trainData = _preprocessorTransform.Transform(_trainData);
+                _validationData = _preprocessorTransform.Transform(_validationData);
             }
 
             var stopwatch = Stopwatch.StartNew();
@@ -97,12 +101,6 @@ public List<RunResult<T>> Execute()
                     // evaluate pipeline
                     runResult = ProcessPipeline(pipeline);
 
-                    if (_preFeaturizers != null)
-                    {
-                        runResult.Estimator = _preFeaturizers.Append(runResult.Estimator);
-                        runResult.Model = preprocessorTransform.Append(runResult.Model);
-                    }
-
                     runResult.RuntimeInSeconds = (int)iterationStopwatch.Elapsed.TotalSeconds;
                     runResult.PipelineInferenceTimeInSeconds = (int)getPiplelineStopwatch.Elapsed.TotalSeconds;
                 }
@@ -129,6 +127,33 @@ public List<RunResult<T>> Execute()
             return iterationResults;
         }
 
+        private static DirectoryInfo GetModelDirectory(DirectoryInfo rootDir)
+        {
+            if (rootDir == null)
+            {
+                return null;
+            }
+            var subdirs = rootDir.Exists ?
+                new HashSet<string>(rootDir.EnumerateDirectories().Select(d => d.Name)) :
+                new HashSet<string>();
+            string experimentDir;
+            for (var i = 0; ; i++)
+            {
+                experimentDir = $"experiment{i}";
+                if (!subdirs.Contains(experimentDir))
+                {
+                    break;
+                }
+            }
+            var experimentDirFullPath = Path.Combine(rootDir.FullName, experimentDir);
+            var experimentDirInfo = new DirectoryInfo(experimentDirFullPath);
+            if (!experimentDirInfo.Exists)
+            {
+                experimentDirInfo.Create();
+            }
+            return experimentDirInfo;
+        }
+
         private void ReportProgress(RunResult<T> iterationResult)
         {
             try
@@ -141,6 +166,17 @@ private void ReportProgress(RunResult<T> iterationResult)
             }
         }
 
+        private FileInfo GetNextModelFileInfo()
+        {
+            if (_experimentSettings.ModelDirectory == null)
+            {
+                return null;
+            }
+
+            return new FileInfo(Path.Combine(_modelDirectory.FullName, 
+                $"Model{_history.Count + 1}.zip"));
+        }
+
         private SuggestedPipelineResult<T> ProcessPipeline(SuggestedPipeline pipeline)
         {
             // run pipeline
@@ -150,22 +186,33 @@ private SuggestedPipelineResult<T> ProcessPipeline(SuggestedPipeline pipeline)
 
             WriteDebugLog(DebugStream.RunResult, $"Processing pipeline {commandLineStr}.");
 
-            var pipelineEstimator = pipeline.ToEstimator();
-
             SuggestedPipelineResult<T> runResult;
 
             try
             {
-                var pipelineModel = pipelineEstimator.Fit(_trainData);
-                var scoredValidationData = pipelineModel.Transform(_validationData);
+                var model = pipeline.ToEstimator().Fit(_trainData);
+                var scoredValidationData = model.Transform(_validationData);
                 var metrics = GetEvaluatedMetrics(scoredValidationData);
                 var score = _metricsAgent.GetScore(metrics);
-                runResult = new SuggestedPipelineResult<T>(metrics, pipelineEstimator, pipelineModel, pipeline, score, null);
+
+                var estimator = pipeline.ToEstimator();
+                if (_preFeaturizers != null)
+                {
+                    estimator = _preFeaturizers.Append(estimator);
+                    model = _preprocessorTransform.Append(model);
+                }
+
+                var modelFileInfo = GetNextModelFileInfo();
+                var modelContainer = modelFileInfo == null ?
+                    new ModelContainer(_context, model) :
+                    new ModelContainer(_context, modelFileInfo, model);
+
+                runResult = new SuggestedPipelineResult<T>(metrics, estimator, modelContainer, pipeline, score, null);
             }
             catch(Exception ex)
             {
                 WriteDebugLog(DebugStream.Exception, $"{pipeline.Trainer} Crashed {ex}");
-                runResult = new SuggestedPipelineResult<T>(null, pipelineEstimator, null, pipeline, 0, ex);
+                runResult = new SuggestedPipelineResult<T>(null, pipeline.ToEstimator(), null, pipeline, 0, ex);
             }
 
             // save pipeline run

src/Microsoft.ML.Auto/Experiment/ModelContainer.cs

@@ -0,0 +1,51 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// 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.IO;
+using Microsoft.ML.Data;
+
+namespace Microsoft.ML.Auto
+{
+    internal class ModelContainer
+    {
+        private readonly MLContext _mlContext;
+        private readonly FileInfo _fileInfo;
+        private readonly ITransformer _model;
+
+        internal ModelContainer(MLContext mlContext, ITransformer model)
+        {
+            _mlContext = mlContext;
+            _model = model;
+        }
+
+        internal ModelContainer(MLContext mlContext, FileInfo fileInfo, ITransformer model)
+        {
+            _mlContext = mlContext;
+            _fileInfo = fileInfo;
+
+            // Write model to disk
+            using (var fs = File.Create(fileInfo.FullName))
+            {
+                model.SaveTo(mlContext, fs);
+            }
+        }
+
+        public ITransformer GetModel()
+        {
+            // If model stored in memory, return it
+            if (_model != null)
+            {
+                return _model;
+            }
+
+            // Load model from disk
+            ITransformer model;
+            using (var stream = new FileStream(_fileInfo.FullName, FileMode.Open, FileAccess.Read, FileShare.Read))
+            {
+                model = _mlContext.Model.Load(stream);
+            }
+            return model;
+        }
+    }
+}

src/Microsoft.ML.Auto/Experiment/SuggestedPipelineResult.cs

@@ -3,6 +3,7 @@
 // See the LICENSE file in the project root for more information.
 
 using System;
+using System.IO;
 
 namespace Microsoft.ML.Auto
 {
@@ -34,25 +35,25 @@ internal class SuggestedPipelineResult<T> : SuggestedPipelineResult
     {
         public readonly T EvaluatedMetrics;
         public IEstimator<ITransformer> Estimator { get; set; }
-        public ITransformer Model { get; set; }
+        public ModelContainer ModelContainer { get; set; }
         public Exception Exception { get; set; }
 
         public int RuntimeInSeconds { get; set; }
         public int PipelineInferenceTimeInSeconds { get; set; }
 
-        public SuggestedPipelineResult(T evaluatedMetrics, IEstimator<ITransformer> estimator, 
-            ITransformer model, SuggestedPipeline pipeline, double score, Exception exception)
+        public SuggestedPipelineResult(T evaluatedMetrics, IEstimator<ITransformer> estimator,
+            ModelContainer modelContainer, SuggestedPipeline pipeline, double score, Exception exception)
             : base(pipeline, score, exception == null)
         {
             EvaluatedMetrics = evaluatedMetrics;
             Estimator = estimator;
-            Model = model;
+            ModelContainer = modelContainer;
             Exception = exception;
         }
 
         public RunResult<T> ToIterationResult()
         {
-            return new RunResult<T>(Model, EvaluatedMetrics, Estimator, Pipeline.ToPipeline(), Exception, RuntimeInSeconds, PipelineInferenceTimeInSeconds);
+            return new RunResult<T>(ModelContainer, EvaluatedMetrics, Estimator, Pipeline.ToPipeline(), Exception, RuntimeInSeconds, PipelineInferenceTimeInSeconds);
         }
     }
 }

src/Test/AutoFitTests.cs

@@ -41,7 +41,7 @@ public void AutoFitMultiTest()
                 .CreateMulticlassClassificationExperiment(0)
                 .Execute(trainData, validationData, new ColumnInformation() { LabelColumn = DatasetUtil.TrivialMulticlassDatasetLabel });
 
-            Assert.IsTrue(result.Max(i => i.ValidationMetrics.AccuracyMacro) > 0.80);
+            Assert.IsTrue(result.Max(i => i.ValidationMetrics.AccuracyMicro) >= 0.8);
         }
 
         [TestMethod]