Add unit tests, and address some code review comments

Author: yaeldMS

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

@@ -251,54 +251,41 @@ private static TransformerChain<ITransformer> Create(IHostEnvironment env, Model
         public static void SaveTo(this ITransformer transformer, IHostEnvironment env, Stream outputStream)
             => new TransformerChain<ITransformer>(transformer).SaveTo(env, outputStream);
 
-        public static ITransformer LoadFrom(IHostEnvironment env, Stream stream)
+        public static ITransformer LoadFromLegacy(IHostEnvironment env, Stream stream)
         {
-            using (var rep = RepositoryReader.Open(stream, env))
-            {
-                try
-                {
-                    ModelLoadContext.LoadModel<ITransformer, SignatureLoadModel>(env, out var transformerChain, rep, LoaderSignature);
-                    return transformerChain;
-                }
-                catch (FormatException ex)
-                {
-                    if (!ex.IsMarked())
-                        throw;
-                    var chain = ModelFileUtils.LoadPipeline(env, stream, new MultiFileSource(null), extractInnerPipe: false);
-                    TransformerChain<ITransformer> transformChain = (chain as LegacyCompositeDataLoader).GetTransformer();
-                    var predictor = ModelFileUtils.LoadPredictorOrNull(env, stream);
-                    if (predictor == null)
-                        return transformChain;
-                    var roles = ModelFileUtils.LoadRoleMappingsOrNull(env, stream);
-                    env.CheckDecode(roles != null, "Predictor model must contain role mappings");
-                    var roleMappings = roles.ToArray();
-
-                    ITransformer pred = null;
-                    if (predictor.PredictionKind == PredictionKind.BinaryClassification)
-                        pred = new BinaryPredictionTransformer<IPredictorProducing<float>>(env, predictor as IPredictorProducing<float>, chain.Schema,
-                            roles.Where(x => x.Key.Value == RoleMappedSchema.ColumnRole.Feature.Value).First().Value);
-                    else if (predictor.PredictionKind == PredictionKind.MulticlassClassification)
-                        pred = new MulticlassPredictionTransformer<IPredictorProducing<VBuffer<float>>>(env,
-                            predictor as IPredictorProducing<VBuffer<float>>, chain.Schema,
-                            roles.Where(x => x.Key.Value == RoleMappedSchema.ColumnRole.Feature.Value).First().Value,
-                            roles.Where(x => x.Key.Value == RoleMappedSchema.ColumnRole.Label.Value).First().Value);
-                    else if (predictor.PredictionKind == PredictionKind.Clustering)
-                        pred = new ClusteringPredictionTransformer<IPredictorProducing<VBuffer<float>>>(env, predictor as IPredictorProducing<VBuffer<float>>, chain.Schema,
-                            roles.Where(x => x.Key.Value == RoleMappedSchema.ColumnRole.Feature.Value).First().Value);
-                    else if (predictor.PredictionKind == PredictionKind.Regression)
-                        pred = new RegressionPredictionTransformer<IPredictorProducing<float>>(env, predictor as IPredictorProducing<float>, chain.Schema,
-                            roles.Where(x => x.Key.Value == RoleMappedSchema.ColumnRole.Feature.Value).First().Value);
-                    else if (predictor.PredictionKind == PredictionKind.AnomalyDetection)
-                        pred = new AnomalyPredictionTransformer<IPredictorProducing<float>>(env, predictor as IPredictorProducing<float>, chain.Schema,
-                            roles.Where(x => x.Key.Value == RoleMappedSchema.ColumnRole.Feature.Value).First().Value);
-                    else if (predictor.PredictionKind == PredictionKind.Ranking)
-                        pred = new RankingPredictionTransformer<IPredictorProducing<float>>(env, predictor as IPredictorProducing<float>, chain.Schema,
-                            roles.Where(x => x.Key.Value == RoleMappedSchema.ColumnRole.Feature.Value).First().Value);
-                    else
-                        throw env.Except("Don't know how to map prediction kind {0}", predictor.PredictionKind);
-                    return transformChain.Append(pred);
-                }
-            }
+            var chain = ModelFileUtils.LoadPipeline(env, stream, new MultiFileSource(null), extractInnerPipe: false);
+            TransformerChain<ITransformer> transformChain = (chain as LegacyCompositeDataLoader).GetTransformer();
+            var predictor = ModelFileUtils.LoadPredictorOrNull(env, stream);
+            if (predictor == null)
+                return transformChain;
+            var roles = ModelFileUtils.LoadRoleMappingsOrNull(env, stream);
+            env.CheckDecode(roles != null, "Predictor model must contain role mappings");
+            var roleMappings = roles.ToArray();
+
+            ITransformer pred = null;
+            if (predictor.PredictionKind == PredictionKind.BinaryClassification)
+                pred = new BinaryPredictionTransformer<IPredictorProducing<float>>(env, predictor as IPredictorProducing<float>, chain.Schema,
+                    roles.Where(x => x.Key.Value == RoleMappedSchema.ColumnRole.Feature.Value).First().Value);
+            else if (predictor.PredictionKind == PredictionKind.MulticlassClassification)
+                pred = new MulticlassPredictionTransformer<IPredictorProducing<VBuffer<float>>>(env,
+                    predictor as IPredictorProducing<VBuffer<float>>, chain.Schema,
+                    roles.Where(x => x.Key.Value == RoleMappedSchema.ColumnRole.Feature.Value).First().Value,
+                    roles.Where(x => x.Key.Value == RoleMappedSchema.ColumnRole.Label.Value).First().Value);
+            else if (predictor.PredictionKind == PredictionKind.Clustering)
+                pred = new ClusteringPredictionTransformer<IPredictorProducing<VBuffer<float>>>(env, predictor as IPredictorProducing<VBuffer<float>>, chain.Schema,
+                    roles.Where(x => x.Key.Value == RoleMappedSchema.ColumnRole.Feature.Value).First().Value);
+            else if (predictor.PredictionKind == PredictionKind.Regression)
+                pred = new RegressionPredictionTransformer<IPredictorProducing<float>>(env, predictor as IPredictorProducing<float>, chain.Schema,
+                    roles.Where(x => x.Key.Value == RoleMappedSchema.ColumnRole.Feature.Value).First().Value);
+            else if (predictor.PredictionKind == PredictionKind.AnomalyDetection)
+                pred = new AnomalyPredictionTransformer<IPredictorProducing<float>>(env, predictor as IPredictorProducing<float>, chain.Schema,
+                    roles.Where(x => x.Key.Value == RoleMappedSchema.ColumnRole.Feature.Value).First().Value);
+            else if (predictor.PredictionKind == PredictionKind.Ranking)
+                pred = new RankingPredictionTransformer<IPredictorProducing<float>>(env, predictor as IPredictorProducing<float>, chain.Schema,
+                    roles.Where(x => x.Key.Value == RoleMappedSchema.ColumnRole.Feature.Value).First().Value);
+            else
+                throw env.Except("Don't know how to map prediction kind {0}", predictor.PredictionKind);
+            return transformChain.Append(pred);
         }
     }
 }

src/Microsoft.ML.Data/Model/ModelOperationsCatalog.cs

@@ -4,6 +4,7 @@
 
 using System;
 using System.IO;
+using System.Linq;
 using Microsoft.Data.DataView;
 using Microsoft.ML.Data;
 using Microsoft.ML.Data.IO;
@@ -32,41 +33,44 @@ internal ModelOperationsCatalog(IHostEnvironment env)
             Explainability = new ExplainabilityTransforms(this);
         }
 
-        private void Save<TSource>(DataViewSchema schema, IDataLoader<TSource> model, Stream stream)
+        /// <summary>
+        /// Save the model to the stream.
+        /// </summary>
+        /// <param name="model">The trained model to be saved.</param>
+        /// <param name="stream">A writeable, seekable stream to save to.</param>
+        public void Save<TSource>(IDataLoader<TSource> model, Stream stream)
         {
+            _env.CheckValue(model, nameof(model));
+            _env.CheckValue(stream, nameof(stream));
+
             using (var rep = RepositoryWriter.CreateNew(stream))
             {
                 ModelSaveContext.SaveModel(rep, model, null);
-                SaveInputSchema(schema, rep);
                 rep.Commit();
             }
         }
 
-        /// <summary>
-        /// Save the model to the stream.
-        /// </summary>
-        /// <param name="model">The trained model to be saved.</param>
-        /// <param name="stream">A writeable, seekable stream to save to.</param>
-        public void Save<TSource>(IDataLoader<TSource> model, Stream stream)
-            => Save(model.GetOutputSchema(), model, stream);
-
         /// <summary>
         /// Save a transformer model and the loader used to create its input data to the stream.
         /// </summary>
         /// <param name="loader">The loader that was used to create data to train the model</param>
         /// <param name="model">The trained model to be saved</param>
         /// <param name="stream">A writeable, seekable stream to save to.</param>
         public void Save<TSource>(IDataLoader<TSource> loader, ITransformer model, Stream stream) =>
-            Save(loader.GetOutputSchema(), new CompositeDataLoader<TSource, ITransformer>(loader, new TransformerChain<ITransformer>(model)), stream);
+            Save(new CompositeDataLoader<TSource, ITransformer>(loader, new TransformerChain<ITransformer>(model)), stream);
 
         /// <summary>
         /// Save a transformer model and the schema of the data that was used to train it to the stream.
         /// </summary>
-        /// <param name="inputSchema">The schema of the input to the transformer.</param>
         /// <param name="model">The trained model to be saved.</param>
+        /// <param name="inputSchema">The schema of the input to the transformer. This can be null.</param>
         /// <param name="stream">A writeable, seekable stream to save to.</param>
-        public void Save(DataViewSchema inputSchema, ITransformer model, Stream stream)
+        public void Save(ITransformer model, DataViewSchema inputSchema, Stream stream)
         {
+            _env.CheckValue(model, nameof(model));
+            _env.CheckValueOrNull(inputSchema);
+            _env.CheckValue(stream, nameof(stream));
+
             using (var rep = RepositoryWriter.CreateNew(stream))
             {
                 ModelSaveContext.SaveModel(rep, model, CompositeDataLoader<object, ITransformer>.TransformerDirectory);
@@ -77,6 +81,12 @@ public void Save(DataViewSchema inputSchema, ITransformer model, Stream stream)
 
         private void SaveInputSchema(DataViewSchema inputSchema, RepositoryWriter rep)
         {
+            _env.AssertValueOrNull(inputSchema);
+            _env.AssertValue(rep);
+
+            if (inputSchema == null)
+                return;
+
             using (var ch = _env.Start("Saving Schema"))
             {
                 var entry = rep.CreateEntry(SchemaEntryName);
@@ -94,30 +104,50 @@ private void SaveInputSchema(DataViewSchema inputSchema, RepositoryWriter rep)
         /// <returns>The loaded model.</returns>
         public ITransformer Load(Stream stream, out DataViewSchema inputSchema)
         {
+            _env.CheckValue(stream, nameof(stream));
+
             using (var rep = RepositoryReader.Open(stream, _env))
             {
                 var entry = rep.OpenEntryOrNull(SchemaEntryName);
                 if (entry != null)
                 {
                     var loader = new BinaryLoader(_env, new BinaryLoader.Arguments(), entry.Stream);
                     inputSchema = loader.Schema;
+                    ModelLoadContext.LoadModel<ITransformer, SignatureLoadModel>(_env, out var transformerChain, rep,
+                        CompositeDataLoader<object, ITransformer>.TransformerDirectory);
+                    return transformerChain;
                 }
-                else
+
+                ModelLoadContext.LoadModelOrNull<IDataLoader<IMultiStreamSource>, SignatureLoadModel>(_env, out var dataLoader, rep, null);
+                if (dataLoader == null)
                 {
+                    // Try to see if the model was saved without a loader or a schema.
+                    if (ModelLoadContext.LoadModelOrNull<ITransformer, SignatureLoadModel>(_env, out var transformerChain, rep,
+                        CompositeDataLoader<object, ITransformer>.TransformerDirectory))
+                    {
+                        inputSchema = null;
+                        return transformerChain;
+                    }
+
                     // Try to load from legacy model format.
                     try
                     {
                         var loader = ModelFileUtils.LoadLoader(_env, rep, new MultiFileSource(null), false);
                         inputSchema = loader.Schema;
+                        return TransformerChain.LoadFromLegacy(_env, stream);
                     }
                     catch (Exception ex)
                     {
-                        if (!ex.IsMarked())
-                            throw;
-                        inputSchema = null;
+                        throw _env.Except(ex, "Could not load legacy format model");
                     }
                 }
-                return TransformerChain.LoadFrom(_env, stream);
+                if (dataLoader is CompositeDataLoader<IMultiStreamSource, ITransformer> composite)
+                {
+                    inputSchema = composite.Loader.GetOutputSchema();
+                    return composite.Transformer;
+                }
+                inputSchema = dataLoader.GetOutputSchema();
+                return new TransformerChain<ITransformer>();
             }
         }
 
@@ -127,12 +157,21 @@ public ITransformer Load(Stream stream, out DataViewSchema inputSchema)
         /// <param name="stream">A readable, seekable stream to load from.</param>
         /// <returns>A model of type <see cref="CompositeDataLoader{IMultiStreamSource, ITransformer}"/> containing the loader
         /// and the transformer chain.</returns>
-        public CompositeDataLoader<IMultiStreamSource, ITransformer> Load(Stream stream)
+        public IDataLoader<IMultiStreamSource> Load(Stream stream)
         {
+            _env.CheckValue(stream, nameof(stream));
+
             using (var rep = RepositoryReader.Open(stream))
             {
-                ModelLoadContext.LoadModel<CompositeDataLoader<IMultiStreamSource, ITransformer>, SignatureLoadModel>(_env, out var model, rep, null);
-                return model;
+                try
+                {
+                    ModelLoadContext.LoadModel<IDataLoader<IMultiStreamSource>, SignatureLoadModel>(_env, out var model, rep, null);
+                    return model;
+                }
+                catch (Exception ex)
+                {
+                    throw _env.Except(ex, "Model does not contain an IDataLoader");
+                }
             }
         }
 
@@ -144,6 +183,8 @@ public CompositeDataLoader<IMultiStreamSource, ITransformer> Load(Stream stream)
         /// <returns>The transformer model from the model stream.</returns>
         public ITransformer Load(Stream stream, out IDataLoader<IMultiStreamSource> loader)
         {
+            _env.CheckValue(stream, nameof(stream));
+
             loader = Load(stream);
             if (loader is CompositeDataLoader<IMultiStreamSource, ITransformer> composite)
             {

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

@@ -5648,7 +5648,6 @@ public void LoadEntryPointModel()
                 {
                     loadedModel = ml.Model.Load(stream, out DataViewSchema inputSchema);
                 }
-
             }
         }
     }