Seperate TreeEnsemble

Author: wschin

src/Microsoft.ML.FastTree/FastTree.cs

@@ -2869,7 +2869,6 @@ internal TreeEnsembleModelParameters(IHostEnvironment env, string name, Internal
             // 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;
-            CreateTreeEnsembleFromInternalDataStructure();
             InnerArgs = innerArgs;
             NumFeatures = numFeatures;
 
@@ -2898,7 +2897,6 @@ protected TreeEnsembleModelParameters(IHostEnvironment env, string name, ModelLo
                 categoricalSplits = true;
 
             TrainedEnsemble = new InternalTreeEnsemble(ctx, usingDefaultValues, categoricalSplits);
-            CreateTreeEnsembleFromInternalDataStructure();
             MaxSplitFeatIdx = TrainedEnsemble.GetMaxFeatureIndex();
 
             InnerArgs = ctx.LoadStringOrNull();
@@ -2919,14 +2917,6 @@ protected TreeEnsembleModelParameters(IHostEnvironment env, string name, ModelLo
             OutputType = NumberType.Float;
         }
 
-        /// <summary>
-        /// This function should be implemented in derived classs to create strongly-typed TreeEnsemble
-        /// from <see cref="TrainedEnsemble"/> and possibly other internal attributes in
-        /// <see cref="TreeEnsembleModelParameters"/>. This also implies we always call this function
-        /// after initializing <see cref="TrainedEnsemble"/>.
-        /// </summary>
-        protected abstract void CreateTreeEnsembleFromInternalDataStructure();
-
         [BestFriend]
         private protected override void SaveCore(ModelSaveContext ctx)
         {
@@ -3400,33 +3390,30 @@ public TreeNode(Dictionary<string, object> keyValues)
     /// </summary>
     public abstract class TreeEnsembleModelParametersBasedOnRegressionTree : TreeEnsembleModelParameters
     {
-        private TreeEnsemble<RegressionTree> _trainedTreeEnsemble;
-
         /// <summary>
         /// An ensemble of trees exposed to users. It is a wrapper on the <see langword="internal"/>
         /// <see cref="InternalTreeEnsemble"/> in <see cref="ML.FastTree.TreeEnsemble{T}"/>.
         /// </summary>
-        public TreeEnsemble<RegressionTree> TrainedTreeEnsemble => _trainedTreeEnsemble;
+        public RegressionTreeEnsemble TrainedTreeEnsemble { get; }
 
         [BestFriend]
         internal TreeEnsembleModelParametersBasedOnRegressionTree(IHostEnvironment env, string name, InternalTreeEnsemble trainedEnsemble, int numFeatures, string innerArgs)
             : base(env, name, trainedEnsemble, numFeatures, innerArgs)
         {
+            TrainedTreeEnsemble = CreateTreeEnsembleFromInternalDataStructure();
         }
 
         protected TreeEnsembleModelParametersBasedOnRegressionTree(IHostEnvironment env, string name, ModelLoadContext ctx, VersionInfo ver)
             : base(env, name, ctx, ver)
         {
+            TrainedTreeEnsemble = CreateTreeEnsembleFromInternalDataStructure();
         }
 
-        /// <summary>
-        /// See <see cref="TreeEnsembleModelParameters.CreateTreeEnsembleFromInternalDataStructure"/>.
-        /// </summary>
-        protected override void CreateTreeEnsembleFromInternalDataStructure()
+        private RegressionTreeEnsemble CreateTreeEnsembleFromInternalDataStructure()
         {
             var trees = TrainedEnsemble.Trees.Select(tree => new RegressionTree(tree));
             var treeWeights = TrainedEnsemble.Trees.Select(tree => tree.Weight);
-            _trainedTreeEnsemble = new TreeEnsemble<RegressionTree>(trees, treeWeights, TrainedEnsemble.Bias);
+            return new RegressionTreeEnsemble(trees, treeWeights, TrainedEnsemble.Bias);
         }
     }
 
@@ -3442,33 +3429,30 @@ protected override void CreateTreeEnsembleFromInternalDataStructure()
     /// </summary>
     public abstract class TreeEnsembleModelParametersBasedOnQuantileRegressionTree : TreeEnsembleModelParameters
     {
-        private TreeEnsemble<QuantileRegressionTree> _trainedTreeEnsemble;
-
         /// <summary>
         /// An ensemble of trees exposed to users. It is a wrapper on the <see langword="internal"/>
         /// <see cref="InternalTreeEnsemble"/> in <see cref="ML.FastTree.TreeEnsemble{T}"/>.
         /// </summary>
-        public TreeEnsemble<QuantileRegressionTree> TrainedTreeEnsemble => _trainedTreeEnsemble;
+        public QuantileRegressionTreeEnsemble TrainedTreeEnsemble { get; }
 
         [BestFriend]
         internal TreeEnsembleModelParametersBasedOnQuantileRegressionTree(IHostEnvironment env, string name, InternalTreeEnsemble trainedEnsemble, int numFeatures, string innerArgs)
             : base(env, name, trainedEnsemble, numFeatures, innerArgs)
         {
+            TrainedTreeEnsemble = CreateTreeEnsembleFromInternalDataStructure();
         }
 
         protected TreeEnsembleModelParametersBasedOnQuantileRegressionTree(IHostEnvironment env, string name, ModelLoadContext ctx, VersionInfo ver)
             : base(env, name, ctx, ver)
         {
+            TrainedTreeEnsemble = CreateTreeEnsembleFromInternalDataStructure();
         }
 
-        /// <summary>
-        /// See <see cref="TreeEnsembleModelParameters.CreateTreeEnsembleFromInternalDataStructure"/>.
-        /// </summary>
-        protected override void CreateTreeEnsembleFromInternalDataStructure()
+        private QuantileRegressionTreeEnsemble CreateTreeEnsembleFromInternalDataStructure()
         {
             var trees = TrainedEnsemble.Trees.Select(tree => new QuantileRegressionTree((InternalQuantileRegressionTree)tree));
             var treeWeights = TrainedEnsemble.Trees.Select(tree => tree.Weight);
-            _trainedTreeEnsemble = new TreeEnsemble<QuantileRegressionTree>(trees, treeWeights, TrainedEnsemble.Bias);
+            return new QuantileRegressionTreeEnsemble(trees, treeWeights, TrainedEnsemble.Bias);
         }
     }
 }

src/Microsoft.ML.FastTree/TreeEnsemble.cs

@@ -12,7 +12,7 @@ namespace Microsoft.ML.FastTree
     /// <see cref="TreeEnsemble{T}"/>, we need to compute the output values of all trees in <see cref="Trees"/>,
     /// scale those values via <see cref="TreeWeights"/>, and finally sum the scaled values and <see cref="Bias"/> up.
     /// </summary>
-    public sealed class TreeEnsemble<T> where T : RegressionTreeBase
+    public abstract class TreeEnsemble<T> where T : RegressionTreeBase
     {
         /// <summary>
         /// When doing prediction, this is a value added to the weighted sum of all <see cref="Trees"/>' outputs.
@@ -37,4 +37,19 @@ internal TreeEnsemble(IEnumerable<T> trees, IEnumerable<double> treeWeights, dou
         }
     }
 
+    public sealed class RegressionTreeEnsemble : TreeEnsemble<RegressionTree>
+    {
+        internal RegressionTreeEnsemble(IEnumerable<RegressionTree> trees, IEnumerable<double> treeWeights, double bias)
+            : base(trees, treeWeights, bias)
+        {
+        }
+    }
+
+    public sealed class QuantileRegressionTreeEnsemble : TreeEnsemble<QuantileRegressionTree>
+    {
+        internal QuantileRegressionTreeEnsemble(IEnumerable<QuantileRegressionTree> trees, IEnumerable<double> treeWeights, double bias)
+            : base(trees, treeWeights, bias)
+        {
+        }
+    }
 }