Typed Calibrated Predictors (#2435)

* Make strongly-typed and weakly-typed accesses possible

Author: wschin

docs/samples/Microsoft.ML.Samples/Dynamic/PermutationFeatureImportance/PfiBinaryClassificationExample.cs

@@ -30,7 +30,7 @@ public static void RunExample()
             var linearPredictor = model.LastTransformer;
             // Linear models for binary classification are wrapped by a calibrator as a generic predictor
             //  To access it directly, we must extract it out and cast it to the proper class
-            var weights = PfiHelper.GetLinearModelWeights(linearPredictor.Model.SubPredictor as LinearBinaryModelParameters);
+            var weights = PfiHelper.GetLinearModelWeights(linearPredictor.Model.SubModel as LinearBinaryModelParameters);
 
             // Compute the permutation metrics using the properly normalized data.
             var transformedData = model.Transform(data);

src/Microsoft.ML.Data/EntryPoints/PredictorModelImpl.cs

@@ -7,6 +7,7 @@
 using System.IO;
 using System.Linq;
 using Microsoft.Data.DataView;
+using Microsoft.ML.Calibrator;
 using Microsoft.ML.Data;
 using Microsoft.ML.Internal.Calibration;
 using Microsoft.ML.Internal.Internallearn;
@@ -111,11 +112,11 @@ internal override string[] GetLabelInfo(IHostEnvironment env, out ColumnType lab
         {
             Contracts.CheckValue(env, nameof(env));
             var predictor = Predictor;
-            var calibrated = predictor as CalibratedPredictorBase;
+            var calibrated = predictor as IWeaklyTypedCalibratedModelParameters;
             while (calibrated != null)
             {
-                predictor = calibrated.SubPredictor;
-                calibrated = predictor as CalibratedPredictorBase;
+                predictor = calibrated.WeeklyTypedSubModel;
+                calibrated = predictor as IWeaklyTypedCalibratedModelParameters;
             }
             var canGetTrainingLabelNames = predictor as ICanGetTrainingLabelNames;
             if (canGetTrainingLabelNames != null)

src/Microsoft.ML.Data/EntryPoints/SummarizePredictor.cs

@@ -5,6 +5,7 @@
 using System.IO;
 using System.Text;
 using Microsoft.Data.DataView;
+using Microsoft.ML.Calibrator;
 using Microsoft.ML.CommandLine;
 using Microsoft.ML.Data;
 using Microsoft.ML.EntryPoints;
@@ -48,11 +49,11 @@ public static CommonOutputs.SummaryOutput Summarize(IHostEnvironment env, Summar
         [BestFriend]
         internal static IDataView GetSummaryAndStats(IHostEnvironment env, IPredictor predictor, RoleMappedSchema schema, out IDataView stats)
         {
-            var calibrated = predictor as CalibratedPredictorBase;
+            var calibrated = predictor as IWeaklyTypedCalibratedModelParameters;
             while (calibrated != null)
             {
-                predictor = calibrated.SubPredictor;
-                calibrated = predictor as CalibratedPredictorBase;
+                predictor = calibrated.WeeklyTypedSubModel;
+                calibrated = predictor as IWeaklyTypedCalibratedModelParameters;
             }
 
             IDataView summary = null;

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

@@ -141,7 +141,6 @@ protected TTransformer TrainTransformer(IDataView trainSet,
             return MakeTransformer(pred, trainSet.Schema);
         }
 
-        [BestFriend]
         private protected abstract TModel TrainModelCore(TrainContext trainContext);
 
         protected abstract TTransformer MakeTransformer(TModel model, Schema trainSchema);

src/Microsoft.ML.Data/Training/TrainerEstimatorBase.cs

@@ -33,14 +33,13 @@ public static SummaryOutput Summarize(IHostEnvironment env, SummarizePredictor.I
 
             input.PredictorModel.PrepareData(host,
                 new EmptyDataView(host, input.PredictorModel.TransformModel.InputSchema),
-                out RoleMappedData rmd, out IPredictor predictor
-);
+                out RoleMappedData rmd, out IPredictor predictor);
 
-            var calibrated = predictor as CalibratedPredictorBase;
+            var calibrated = predictor as IWeaklyTypedCalibratedModelParameters;
             while (calibrated != null)
             {
-                predictor = calibrated.SubPredictor;
-                calibrated = predictor as CalibratedPredictorBase;
+                predictor = calibrated.WeeklyTypedSubModel;
+                calibrated = predictor as IWeaklyTypedCalibratedModelParameters;
             }
             var ensemble = predictor as SchemaBindablePipelineEnsembleBase;
             host.CheckUserArg(ensemble != null, nameof(input.PredictorModel.Predictor), "Predictor is not a pipeline ensemble predictor");

src/Microsoft.ML.Ensemble/EntryPoints/PipelineEnsemble.cs

@@ -96,15 +96,17 @@ private static IPredictorProducing<float> Create(IHostEnvironment env, ModelLoad
             ctx.LoadModelOrNull<ICalibrator, SignatureLoadModel>(env, out calibrator, @"Calibrator");
             if (calibrator == null)
                 return predictor;
-            return new SchemaBindableCalibratedPredictor(env, predictor, calibrator);
+            return new SchemaBindableCalibratedModelParameters<FastTreeBinaryModelParameters, ICalibrator>(env, predictor, calibrator);
         }
 
         public override PredictionKind PredictionKind => PredictionKind.BinaryClassification;
     }
 
     /// <include file = 'doc.xml' path='doc/members/member[@name="FastTree"]/*' />
     public sealed partial class FastTreeBinaryClassificationTrainer :
-        BoostingFastTreeTrainerBase<FastTreeBinaryClassificationTrainer.Options, BinaryPredictionTransformer<IPredictorWithFeatureWeights<float>>, IPredictorWithFeatureWeights<float>>
+        BoostingFastTreeTrainerBase<FastTreeBinaryClassificationTrainer.Options,
+        BinaryPredictionTransformer<CalibratedModelParametersBase<FastTreeBinaryModelParameters, PlattCalibrator>>,
+        CalibratedModelParametersBase<FastTreeBinaryModelParameters, PlattCalibrator>>
     {
         /// <summary>
         /// The LoadName for the assembly containing the trainer.
@@ -156,7 +158,7 @@ internal FastTreeBinaryClassificationTrainer(IHostEnvironment env, Options optio
 
         public override PredictionKind PredictionKind => PredictionKind.BinaryClassification;
 
-        private protected override IPredictorWithFeatureWeights<float> TrainModelCore(TrainContext context)
+        private protected override CalibratedModelParametersBase<FastTreeBinaryModelParameters, PlattCalibrator> TrainModelCore(TrainContext context)
         {
             Host.CheckValue(context, nameof(context));
             var trainData = context.TrainingSet;
@@ -185,7 +187,7 @@ private protected override IPredictorWithFeatureWeights<float> TrainModelCore(Tr
             // BinaryClassificationObjectiveFunction.GetGradientInOneQuery being consistent with the
             // description in section 6 of the paper.
             var cali = new PlattCalibrator(Host, -1 * _sigmoidParameter, 0);
-            return new FeatureWeightsCalibratedPredictor(Host, pred, cali);
+            return new FeatureWeightsCalibratedModelParameters<FastTreeBinaryModelParameters, PlattCalibrator>(Host, pred, cali);
         }
 
         protected override ObjectiveFunctionBase ConstructObjFunc(IChannel ch)
@@ -273,10 +275,11 @@ protected override void InitializeTests()
             }
         }
 
-        protected override BinaryPredictionTransformer<IPredictorWithFeatureWeights<float>> MakeTransformer(IPredictorWithFeatureWeights<float> model, Schema trainSchema)
-        => new BinaryPredictionTransformer<IPredictorWithFeatureWeights<float>>(Host, model, trainSchema, FeatureColumn.Name);
+        protected override BinaryPredictionTransformer<CalibratedModelParametersBase<FastTreeBinaryModelParameters, PlattCalibrator>> MakeTransformer(
+            CalibratedModelParametersBase<FastTreeBinaryModelParameters, PlattCalibrator> model, Schema trainSchema)
+            => new BinaryPredictionTransformer<CalibratedModelParametersBase<FastTreeBinaryModelParameters, PlattCalibrator>>(Host, model, trainSchema, FeatureColumn.Name);
 
-        public BinaryPredictionTransformer<IPredictorWithFeatureWeights<float>> Train(IDataView trainData, IDataView validationData = null)
+        public BinaryPredictionTransformer<CalibratedModelParametersBase<FastTreeBinaryModelParameters, PlattCalibrator>> Train(IDataView trainData, IDataView validationData = null)
             => TrainTransformer(trainData, validationData);
 
         protected override SchemaShape.Column[] GetOutputColumnsCore(SchemaShape inputSchema)

src/Microsoft.ML.FastTree/FastTreeClassification.cs

@@ -31,7 +31,9 @@
 namespace Microsoft.ML.Trainers.FastTree
 {
     public sealed class BinaryClassificationGamTrainer :
-    GamTrainerBase<BinaryClassificationGamTrainer.Options, BinaryPredictionTransformer<CalibratedPredictorBase>, CalibratedPredictorBase>
+        GamTrainerBase<BinaryClassificationGamTrainer.Options,
+        BinaryPredictionTransformer<CalibratedModelParametersBase<BinaryClassificationGamModelParameters, PlattCalibrator>>,
+        CalibratedModelParametersBase<BinaryClassificationGamModelParameters, PlattCalibrator>>
     {
         public sealed class Options : ArgumentsBase
         {
@@ -103,14 +105,13 @@ private static bool[] ConvertTargetsToBool(double[] targets)
             Parallel.Invoke(new ParallelOptions { MaxDegreeOfParallelism = BlockingThreadPool.NumThreads }, actions);
             return boolArray;
         }
-
-        private protected override CalibratedPredictorBase TrainModelCore(TrainContext context)
+        private protected override CalibratedModelParametersBase<BinaryClassificationGamModelParameters, PlattCalibrator> TrainModelCore(TrainContext context)
         {
             TrainBase(context);
             var predictor = new BinaryClassificationGamModelParameters(Host,
                 BinUpperBounds, BinEffects, MeanEffect, InputLength, FeatureMap);
             var calibrator = new PlattCalibrator(Host, -1.0 * _sigmoidParameter, 0);
-            return new CalibratedPredictor(Host, predictor, calibrator);
+            return new ValueMapperCalibratedModelParameters<BinaryClassificationGamModelParameters, PlattCalibrator>(Host, predictor, calibrator);
         }
 
         protected override ObjectiveFunctionBase CreateObjectiveFunction()
@@ -139,10 +140,11 @@ protected override void DefinePruningTest()
             PruningTest = new TestHistory(validTest, PruningLossIndex);
         }
 
-        protected override BinaryPredictionTransformer<CalibratedPredictorBase> MakeTransformer(CalibratedPredictorBase model, Schema trainSchema)
-            => new BinaryPredictionTransformer<CalibratedPredictorBase>(Host, model, trainSchema, FeatureColumn.Name);
+        protected override BinaryPredictionTransformer<CalibratedModelParametersBase<BinaryClassificationGamModelParameters, PlattCalibrator>>
+            MakeTransformer(CalibratedModelParametersBase<BinaryClassificationGamModelParameters, PlattCalibrator> model, Schema trainSchema)
+            => new BinaryPredictionTransformer<CalibratedModelParametersBase<BinaryClassificationGamModelParameters, PlattCalibrator>>(Host, model, trainSchema, FeatureColumn.Name);
 
-        public BinaryPredictionTransformer<CalibratedPredictorBase> Train(IDataView trainData, IDataView validationData = null)
+        public BinaryPredictionTransformer<CalibratedModelParametersBase<BinaryClassificationGamModelParameters, PlattCalibrator>> Train(IDataView trainData, IDataView validationData = null)
             => TrainTransformer(trainData, validationData);
 
         protected override SchemaShape.Column[] GetOutputColumnsCore(SchemaShape inputSchema)
@@ -207,7 +209,7 @@ private static IPredictorProducing<float> Create(IHostEnvironment env, ModelLoad
             ctx.LoadModelOrNull<ICalibrator, SignatureLoadModel>(env, out calibrator, @"Calibrator");
             if (calibrator == null)
                 return predictor;
-            return new SchemaBindableCalibratedPredictor(env, predictor, calibrator);
+            return new SchemaBindableCalibratedModelParameters<BinaryClassificationGamModelParameters, ICalibrator>(env, predictor, calibrator);
         }
 
         private protected override void SaveCore(ModelSaveContext ctx)

src/Microsoft.ML.FastTree/GamClassification.cs

@@ -876,12 +876,17 @@ private Context Init(IChannel ch)
                 LoadModelObjects(ch, true, out rawPred, true, out schema, out loader);
                 bool hadCalibrator = false;
 
-                var calibrated = rawPred as CalibratedPredictorBase;
+                // The rawPred has two possible types:
+                //  1. CalibratedPredictorBase<BinaryClassificationGamModelParameters, PlattCalibrator>
+                //  2. RegressionGamModelParameters
+                // For (1), the trained model, GamModelParametersBase, is a field we need to extract. For (2),
+                // we don't need to do anything because RegressionGamModelParameters is derived from GamModelParametersBase.
+                var calibrated = rawPred as CalibratedModelParametersBase<BinaryClassificationGamModelParameters, PlattCalibrator>;
                 while (calibrated != null)
                 {
                     hadCalibrator = true;
-                    rawPred = calibrated.SubPredictor;
-                    calibrated = rawPred as CalibratedPredictorBase;
+                    rawPred = calibrated.SubModel;
+                    calibrated = rawPred as CalibratedModelParametersBase<BinaryClassificationGamModelParameters, PlattCalibrator>;
                 }
                 var pred = rawPred as GamModelParametersBase;
                 ch.CheckUserArg(pred != null, nameof(Args.InputModelFile), "Predictor was not a " + nameof(GamModelParametersBase));

src/Microsoft.ML.FastTree/GamModelParameters.cs

@@ -105,13 +105,13 @@ private static IPredictorProducing<float> Create(IHostEnvironment env, ModelLoad
             ctx.LoadModelOrNull<ICalibrator, SignatureLoadModel>(env, out calibrator, @"Calibrator");
             if (calibrator == null)
                 return predictor;
-            return new SchemaBindableCalibratedPredictor(env, predictor, calibrator);
+            return new SchemaBindableCalibratedModelParameters<FastForestClassificationModelParameters, ICalibrator>(env, predictor, calibrator);
         }
     }
 
     /// <include file='doc.xml' path='doc/members/member[@name="FastForest"]/*' />
     public sealed partial class FastForestClassification :
-        RandomForestTrainerBase<FastForestClassification.Options, BinaryPredictionTransformer<IPredictorWithFeatureWeights<float>>, IPredictorWithFeatureWeights<float>>
+        RandomForestTrainerBase<FastForestClassification.Options, BinaryPredictionTransformer<FastForestClassificationModelParameters>, FastForestClassificationModelParameters>
     {
         public sealed class Options : FastForestArgumentsBase
         {
@@ -170,7 +170,7 @@ internal FastForestClassification(IHostEnvironment env, Options options)
         {
         }
 
-        private protected override IPredictorWithFeatureWeights<float> TrainModelCore(TrainContext context)
+        private protected override FastForestClassificationModelParameters TrainModelCore(TrainContext context)
         {
             Host.CheckValue(context, nameof(context));
             var trainData = context.TrainingSet;
@@ -213,10 +213,10 @@ protected override Test ConstructTestForTrainingData()
             return new BinaryClassificationTest(ConstructScoreTracker(TrainSet), _trainSetLabels, 1);
         }
 
-        protected override BinaryPredictionTransformer<IPredictorWithFeatureWeights<float>> MakeTransformer(IPredictorWithFeatureWeights<float> model, Schema trainSchema)
-         => new BinaryPredictionTransformer<IPredictorWithFeatureWeights<float>>(Host, model, trainSchema, FeatureColumn.Name);
+        protected override BinaryPredictionTransformer<FastForestClassificationModelParameters> MakeTransformer(FastForestClassificationModelParameters model, Schema trainSchema)
+         => new BinaryPredictionTransformer<FastForestClassificationModelParameters>(Host, model, trainSchema, FeatureColumn.Name);
 
-        public BinaryPredictionTransformer<IPredictorWithFeatureWeights<float>> Train(IDataView trainData, IDataView validationData = null)
+        public BinaryPredictionTransformer<FastForestClassificationModelParameters> Train(IDataView trainData, IDataView validationData = null)
             => TrainTransformer(trainData, validationData);
 
         protected override SchemaShape.Column[] GetOutputColumnsCore(SchemaShape inputSchema)