Immutable RegressionTree and TreeEnsemble wrappers

Author: wschin

src/Microsoft.ML.FastTree/FastTree.cs

@@ -2809,9 +2809,10 @@ public abstract class TreeEnsembleModelParameters :
         ISingleCanSavePfa,
         ISingleCanSaveOnnx
     {
+        public TreeEnsembleView TrainedTreeEnsembleView { get; }
         //The below two properties are necessary for tree Visualizer
         [BestFriend]
-        internal TreeEnsemble TrainedEnsemble { get; }
+        internal TreeEnsemble TrainedEnsemble => TrainedTreeEnsembleView.TreeEnsemble;
         int ITreeEnsemble.NumTrees => TrainedEnsemble.NumTrees;
 
         // Inner args is used only for documentation purposes when saving comments to INI files.
@@ -2853,7 +2854,29 @@ public abstract class TreeEnsembleModelParameters :
         /// </summary>
         public FeatureContributionCalculator FeatureContributionCalculator => new FeatureContributionCalculator(this);
 
-        public TreeEnsembleModelParameters(IHostEnvironment env, string name, TreeEnsemble trainedEnsemble, int numFeatures, string innerArgs)
+        public TreeEnsembleModelParameters(IHostEnvironment env, string name, TreeEnsembleView trainedEnsembleView,
+            int numFeatures)
+            : base(env, name)
+        {
+            Host.CheckValue(trainedEnsembleView, nameof(trainedEnsembleView));
+            Host.CheckParam(numFeatures > 0, nameof(numFeatures), "must be positive");
+
+            // REVIEW: When we make the predictor wrapper, we may want to further "optimize"
+            // the trained ensemble to, for instance, resize arrays so that they are of the length
+            // the actual number of leaves/nodes, or remove unnecessary arrays, and so forth.
+            TrainedTreeEnsembleView = trainedEnsembleView;
+            InnerArgs = "";
+            NumFeatures = trainedEnsembleView.Trees.Select(tree => tree.ActiveFeatures.Length).Max();
+
+            MaxSplitFeatIdx = trainedEnsembleView.TreeEnsemble.GetMaxFeatureIndex();
+            Contracts.Assert(NumFeatures > MaxSplitFeatIdx);
+
+            InputType = new VectorType(NumberType.Float, NumFeatures);
+            OutputType = NumberType.Float;
+        }
+
+        [BestFriend]
+        internal TreeEnsembleModelParameters(IHostEnvironment env, string name, TreeEnsemble trainedEnsemble, int numFeatures, string innerArgs)
             : base(env, name)
         {
             Host.CheckValue(trainedEnsemble, nameof(trainedEnsemble));
@@ -2863,11 +2886,11 @@ public TreeEnsembleModelParameters(IHostEnvironment env, string name, TreeEnsemb
             // REVIEW: When we make the predictor wrapper, we may want to further "optimize"
             // the trained ensemble to, for instance, resize arrays so that they are of the length
             // the actual number of leaves/nodes, or remove unnecessary arrays, and so forth.
-            TrainedEnsemble = trainedEnsemble;
+            TrainedTreeEnsembleView = new TreeEnsembleView(trainedEnsemble);
             InnerArgs = innerArgs;
             NumFeatures = numFeatures;
 
-            MaxSplitFeatIdx = FindMaxFeatureIndex(trainedEnsemble);
+            MaxSplitFeatIdx = trainedEnsemble.GetMaxFeatureIndex();
             Contracts.Assert(NumFeatures > MaxSplitFeatIdx);
 
             InputType = new VectorType(NumberType.Float, NumFeatures);
@@ -2891,8 +2914,8 @@ protected TreeEnsembleModelParameters(IHostEnvironment env, string name, ModelLo
             if (ctx.Header.ModelVerWritten >= VerCategoricalSplitSerialized)
                 categoricalSplits = true;
 
-            TrainedEnsemble = new TreeEnsemble(ctx, usingDefaultValues, categoricalSplits);
-            MaxSplitFeatIdx = FindMaxFeatureIndex(TrainedEnsemble);
+            TrainedTreeEnsembleView = new TreeEnsembleView(new TreeEnsemble(ctx, usingDefaultValues, categoricalSplits));
+            MaxSplitFeatIdx = TrainedEnsemble.GetMaxFeatureIndex();
 
             InnerArgs = ctx.LoadStringOrNull();
             if (ctx.Header.ModelVerWritten >= VerNumFeaturesSerialized)
@@ -3258,23 +3281,6 @@ public void GetFeatureWeights(ref VBuffer<Float> weights)
             bldr.GetResult(ref weights);
         }
 
-        private static int FindMaxFeatureIndex(TreeEnsemble ensemble)
-        {
-            int ifeatMax = 0;
-            for (int i = 0; i < ensemble.NumTrees; i++)
-            {
-                var tree = ensemble.GetTreeAt(i);
-                for (int n = 0; n < tree.NumNodes; n++)
-                {
-                    int ifeat = tree.SplitFeature(n);
-                    if (ifeat > ifeatMax)
-                        ifeatMax = ifeat;
-                }
-            }
-
-            return ifeatMax;
-        }
-
         ITree[] ITreeEnsemble.GetTrees()
         {
             return TrainedEnsemble.Trees.Select(k => new Tree(k)).ToArray();