Make IDataLoader<TSource> inherit from ICanSaveModel, and add methods to save/load it in ModelOperationsCatalog.

Author: yaeldMS

src/Microsoft.ML.Core/Data/IEstimator.cs

@@ -224,7 +224,7 @@ internal bool TryFindColumn(string name, out Column column)
     /// The 'data loader' takes a certain kind of input and turns it into an <see cref="IDataView"/>.
     /// </summary>
     /// <typeparam name="TSource">The type of input the loader takes.</typeparam>
-    public interface IDataLoader<in TSource>
+    public interface IDataLoader<in TSource> : ICanSaveModel
     {
         /// <summary>
         /// Produce the data view from the specified input.

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

@@ -4,7 +4,11 @@
 
 using System.IO;
 using Microsoft.Data.DataView;
-using Microsoft.ML.Model;
+using Microsoft.ML;
+using Microsoft.ML.Data;
+
+[assembly: LoadableClass(CompositeDataLoader<IMultiStreamSource, ITransformer>.Summary, typeof(CompositeDataLoader<IMultiStreamSource, ITransformer>), null, typeof(SignatureLoadModel),
+    "Composite Loader", CompositeDataLoader<IMultiStreamSource, ITransformer>.LoaderSignature)]
 
 namespace Microsoft.ML.Data
 {
@@ -15,6 +19,10 @@ namespace Microsoft.ML.Data
     public sealed class CompositeDataLoader<TSource, TLastTransformer> : IDataLoader<TSource>
         where TLastTransformer : class, ITransformer
     {
+        private const string LoaderDirectory = "Loader";
+        private const string LegacyLoaderDirectory = "Reader";
+        private const string TransformerDirectory = TransformerChain.LoaderSignature;
+
         /// <summary>
         /// The underlying data loader.
         /// </summary>
@@ -33,6 +41,24 @@ public CompositeDataLoader(IDataLoader<TSource> loader, TransformerChain<TLastTr
             Transformer = transformerChain ?? new TransformerChain<TLastTransformer>();
         }
 
+        private CompositeDataLoader(IHost host, ModelLoadContext ctx)
+        {
+            if (!ctx.LoadModelOrNull<IDataLoader<TSource>, SignatureLoadModel>(host, out Loader, LegacyLoaderDirectory))
+                ctx.LoadModel<IDataLoader<TSource>, SignatureLoadModel>(host, out Loader, LoaderDirectory);
+            ctx.LoadModel<TransformerChain<TLastTransformer>, SignatureLoadModel>(host, out Transformer, TransformerDirectory);
+        }
+
+        private static CompositeDataLoader<TSource, TLastTransformer> Create(IHostEnvironment env, ModelLoadContext ctx)
+        {
+            Contracts.CheckValue(env, nameof(env));
+            IHost h = env.Register(LoaderSignature);
+
+            h.CheckValue(ctx, nameof(ctx));
+            ctx.CheckAtModel(GetVersionInfo());
+
+            return h.Apply("Loading Model", ch => new CompositeDataLoader<TSource, TLastTransformer>(h, ctx));
+        }
+
         /// <summary>
         /// Produce the data view from the specified input.
         /// Note that <see cref="IDataView"/>'s are lazy, so no actual loading happens here, just schema validation.
@@ -62,6 +88,16 @@ public CompositeDataLoader<TSource, TNewLast> AppendTransformer<TNewLast>(TNewLa
             return new CompositeDataLoader<TSource, TNewLast>(Loader, Transformer.Append(transformer));
         }
 
+        void ICanSaveModel.Save(ModelSaveContext ctx)
+        {
+            Contracts.CheckValue(ctx, nameof(ctx));
+            ctx.CheckAtModel();
+            ctx.SetVersionInfo(GetVersionInfo());
+
+            ctx.SaveModel(Loader, LoaderDirectory);
+            ctx.SaveModel(Transformer, TransformerDirectory);
+        }
+
         /// <summary>
         /// Save the contents to a stream, as a "model file".
         /// </summary>
@@ -76,47 +112,27 @@ public void SaveTo(IHostEnvironment env, Stream outputStream)
                 using (var rep = RepositoryWriter.CreateNew(outputStream, ch))
                 {
                     ch.Trace("Saving data loader");
-                    ModelSaveContext.SaveModel(rep, Loader, "Reader");
+                    ModelSaveContext.SaveModel(rep, Loader, LoaderDirectory);
 
                     ch.Trace("Saving transformer chain");
-                    ModelSaveContext.SaveModel(rep, Transformer, TransformerChain.LoaderSignature);
+                    ModelSaveContext.SaveModel(rep, Transformer, TransformerDirectory);
                     rep.Commit();
                 }
             }
         }
-    }
 
-    /// <summary>
-    /// Utility class to facilitate loading from a stream.
-    /// </summary>
-    [BestFriend]
-    internal static class CompositeDataLoader
-    {
-        /// <summary>
-        /// Save the contents to a stream, as a "model file".
-        /// </summary>
-        public static void SaveTo<TSource>(this IDataLoader<TSource> loader, IHostEnvironment env, Stream outputStream)
-            => new CompositeDataLoader<TSource, ITransformer>(loader).SaveTo(env, outputStream);
+        internal const string Summary = "A loader that encapsulates a loader and a transformer chain.";
 
-        /// <summary>
-        /// Load the pipeline from stream.
-        /// </summary>
-        public static CompositeDataLoader<IMultiStreamSource, ITransformer> LoadFrom(IHostEnvironment env, Stream stream)
+        internal const string LoaderSignature = "CompositeLoader";
+        private static VersionInfo GetVersionInfo()
         {
-            Contracts.CheckValue(env, nameof(env));
-            env.CheckValue(stream, nameof(stream));
-
-            env.Check(stream.CanRead && stream.CanSeek, "Need a readable and seekable stream to load");
-            using (var rep = RepositoryReader.Open(stream, env))
-            using (var ch = env.Start("Loading pipeline"))
-            {
-                ch.Trace("Loading data loader");
-                ModelLoadContext.LoadModel<IDataLoader<IMultiStreamSource>, SignatureLoadModel>(env, out var loader, rep, "Reader");
-
-                ch.Trace("Loader transformer chain");
-                ModelLoadContext.LoadModel<TransformerChain<ITransformer>, SignatureLoadModel>(env, out var transformerChain, rep, TransformerChain.LoaderSignature);
-                return new CompositeDataLoader<IMultiStreamSource, ITransformer>(loader, transformerChain);
-            }
+            return new VersionInfo(
+                modelSignature: "CMPSTLDR",
+                verWrittenCur: 0x00010001, // Initial
+                verReadableCur: 0x00010001,
+                verWeCanReadBack: 0x00010001,
+                loaderSignature: LoaderSignature,
+                loaderAssemblyName: typeof(CompositeDataLoader<,>).Assembly.FullName);
         }
     }
 }

src/Microsoft.ML.Data/DataLoadSave/Text/TextLoader.cs

@@ -12,20 +12,22 @@
 using Microsoft.ML.CommandLine;
 using Microsoft.ML.Data;
 using Microsoft.ML.Internal.Utilities;
-using Microsoft.ML.Model;
 
 [assembly: LoadableClass(TextLoader.Summary, typeof(ILegacyDataLoader), typeof(TextLoader), typeof(TextLoader.Options), typeof(SignatureDataLoader),
     "Text Loader", "TextLoader", "Text", DocName = "loader/TextLoader.md")]
 
 [assembly: LoadableClass(TextLoader.Summary, typeof(ILegacyDataLoader), typeof(TextLoader), null, typeof(SignatureLoadDataLoader),
     "Text Loader", TextLoader.LoaderSignature)]
 
+[assembly: LoadableClass(TextLoader.Summary, typeof(TextLoader), null, typeof(SignatureLoadModel),
+    "Text Loader", TextLoader.LoaderSignature)]
+
 namespace Microsoft.ML.Data
 {
     /// <summary>
     /// Loads a text file into an IDataView. Supports basic mapping from input columns to <see cref="IDataView"/> columns.
     /// </summary>
-    public sealed partial class TextLoader : IDataLoader<IMultiStreamSource>, ICanSaveModel
+    public sealed partial class TextLoader : IDataLoader<IMultiStreamSource>
     {
         /// <summary>
         /// Describes how an input column should be mapped to an <see cref="IDataView"/> column.

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

@@ -255,7 +255,9 @@ public static TransformerChain<ITransformer> LoadFrom(IHostEnvironment env, Stre
             {
                 try
                 {
-                    ModelLoadContext.LoadModel<TransformerChain<ITransformer>, SignatureLoadModel>(env, out var transformerChain, rep, LoaderSignature);
+                    ModelLoadContext.LoadModelOrNull<TransformerChain<ITransformer>, SignatureLoadModel>(env, out var transformerChain, rep, LoaderSignature);
+                    if (transformerChain == null)
+                        ModelLoadContext.LoadModel<TransformerChain<ITransformer>, SignatureLoadModel>(env, out transformerChain, rep, $@"Model\{LoaderSignature}");
                     return transformerChain;
                 }
                 catch (FormatException ex)

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

@@ -45,13 +45,34 @@ protected SubCatalogBase(ModelOperationsCatalog owner)
         /// <param name="stream">A writeable, seekable stream to save to.</param>
         public void Save(ITransformer model, Stream stream) => model.SaveTo(Environment, stream);
 
+        public void Save<TSource>(IDataLoader<TSource> model, Stream stream)
+        {
+            using (var rep = RepositoryWriter.CreateNew(stream))
+            {
+                ModelSaveContext.SaveModel(rep, model, "Model");
+                rep.Commit();
+            }
+        }
+
+        public void Save<TSource>(IDataLoader<TSource> loader, ITransformer model, Stream stream) =>
+            Save(new CompositeDataLoader<TSource, ITransformer>(loader, new TransformerChain<ITransformer>(model)), stream);
+
         /// <summary>
         /// Load the model from the stream.
         /// </summary>
         /// <param name="stream">A readable, seekable stream to load from.</param>
         /// <returns>The loaded model.</returns>
         public ITransformer Load(Stream stream) => TransformerChain.LoadFrom(Environment, stream);
 
+        public IDataLoader<IMultiStreamSource> LoadAsCompositeDataLoader(Stream stream)
+        {
+            using (var rep = RepositoryReader.Open(stream))
+            {
+                ModelLoadContext.LoadModel<IDataLoader<IMultiStreamSource>, SignatureLoadModel>(Environment, out var model, rep, "Model");
+                return model;
+            }
+        }
+
         /// <summary>
         /// The catalog of model explainability operations.
         /// </summary>

test/Microsoft.ML.Tests/Scenarios/Api/Estimators/DeserializationTests.cs

@@ -53,5 +53,62 @@ public void LoadModelAndExtractPredictor()
                 as BinaryClassificationGamModelParameters;
             Assert.NotNull(gam);
         }
+
+        [Fact]
+        public void SaveAndLoadModelWithLoader()
+        {
+            var ml = new MLContext(seed: 1, conc: 1);
+            var file = new MultiFileSource(GetDataPath(TestDatasets.adult.trainFilename));
+            var loader = ml.Data.CreateTextLoader<InputData>(hasHeader: true, dataSample: file);
+            var data = loader.Load(file);
+
+            // Pipeline.
+            var pipeline = ml.BinaryClassification.Trainers.GeneralizedAdditiveModels();
+
+            // Train.
+            var model = pipeline.Fit(data);
+
+            // Save and reload.
+            string modelPath = GetOutputPath(FullTestName + "-model.zip");
+            using (var fs = File.Create(modelPath))
+                ml.Model.Save(loader, model, fs);
+
+            IDataLoader<IMultiStreamSource> loadedModel;
+            ITransformer loadedModelWithoutLoader;
+            using (var fs = File.OpenRead(modelPath))
+            {
+                loadedModel = ml.Model.LoadAsCompositeDataLoader(fs);
+                loadedModelWithoutLoader = ml.Model.Load(fs);
+            }
+
+            // Without deserializing the loader from the model we lose the slot names.
+            data = ml.Data.LoadFromEnumerable(new[] { new InputData() });
+            data = loadedModelWithoutLoader.Transform(data);
+            Assert.Null(data.Schema["Features"].Annotations.Schema.GetColumnOrNull(AnnotationUtils.Kinds.SlotNames));
+
+            data = loadedModel.Load(file);
+            Assert.True(data.Schema["Features"].HasSlotNames(data.Schema["Features"].Type.GetValueCount()));
+            VBuffer<ReadOnlyMemory<char>> slotNames = default;
+            data.Schema["Features"].GetSlotNames(ref slotNames);
+            var ageIndex = FindIndex(slotNames.GetValues(), "age");
+            var transformer = (loadedModel as CompositeDataLoader<IMultiStreamSource, ITransformer>).Transformer.LastTransformer;
+            var gamModel = ((transformer as BinaryPredictionTransformer<object>).Model
+                as CalibratedModelParametersBase<object, ICalibrator>).SubModel
+                as BinaryClassificationGamModelParameters;
+            var ageBinUpperBounds = gamModel.GetBinUpperBounds(ageIndex);
+            var ageBinEffects = gamModel.GetBinEffects(ageIndex);
+        }
+
+        private int FindIndex(ReadOnlySpan<ReadOnlyMemory<char>> values, string slotName)
+        {
+            int index = 0;
+            foreach (var value in values)
+            {
+                if (value.Span.SequenceEqual(slotName.AsSpan()))
+                    return index;
+                index++;
+            }
+            return -1;
+        }
     }
 }