| # |
| # Licensed to the Apache Software Foundation (ASF) under one or more |
| # contributor license agreements. See the NOTICE file distributed with |
| # this work for additional information regarding copyright ownership. |
| # The ASF licenses this file to You under the Apache License, Version 2.0 |
| # (the "License"); you may not use this file except in compliance with |
| # the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| # |
| |
| import sys |
| import random |
| |
| from pyspark import RDD, since |
| from pyspark.mllib.common import callMLlibFunc, inherit_doc, JavaModelWrapper |
| from pyspark.mllib.linalg import _convert_to_vector |
| from pyspark.mllib.regression import LabeledPoint |
| from pyspark.mllib.util import JavaLoader, JavaSaveable |
| from typing import Dict, Optional, Tuple, Union, overload, TYPE_CHECKING |
| from pyspark.core.rdd import RDD |
| |
| if TYPE_CHECKING: |
| from pyspark.mllib._typing import VectorLike |
| |
| |
| __all__ = [ |
| "DecisionTreeModel", |
| "DecisionTree", |
| "RandomForestModel", |
| "RandomForest", |
| "GradientBoostedTreesModel", |
| "GradientBoostedTrees", |
| ] |
| |
| |
| class TreeEnsembleModel(JavaModelWrapper, JavaSaveable): |
| """TreeEnsembleModel |
| |
| .. versionadded:: 1.3.0 |
| """ |
| |
| @overload |
| def predict(self, x: "VectorLike") -> float: |
| ... |
| |
| @overload |
| def predict(self, x: RDD["VectorLike"]) -> RDD[float]: |
| ... |
| |
| def predict(self, x: Union["VectorLike", RDD["VectorLike"]]) -> Union[float, RDD[float]]: |
| """ |
| Predict values for a single data point or an RDD of points using |
| the model trained. |
| |
| .. versionadded:: 1.3.0 |
| |
| Notes |
| ----- |
| In Python, predict cannot currently be used within an RDD |
| transformation or action. |
| Call predict directly on the RDD instead. |
| """ |
| if isinstance(x, RDD): |
| return self.call("predict", x.map(_convert_to_vector)) |
| |
| else: |
| return self.call("predict", _convert_to_vector(x)) |
| |
| @since("1.3.0") |
| def numTrees(self) -> int: |
| """ |
| Get number of trees in ensemble. |
| """ |
| return self.call("numTrees") |
| |
| @since("1.3.0") |
| def totalNumNodes(self) -> int: |
| """ |
| Get total number of nodes, summed over all trees in the ensemble. |
| """ |
| return self.call("totalNumNodes") |
| |
| def __repr__(self) -> str: |
| """Summary of model""" |
| return self._java_model.toString() |
| |
| @since("1.3.0") |
| def toDebugString(self) -> str: |
| """Full model""" |
| return self._java_model.toDebugString() |
| |
| |
| class DecisionTreeModel(JavaModelWrapper, JavaSaveable, JavaLoader["DecisionTreeModel"]): |
| """ |
| A decision tree model for classification or regression. |
| |
| .. versionadded:: 1.1.0 |
| """ |
| |
| @overload |
| def predict(self, x: "VectorLike") -> float: |
| ... |
| |
| @overload |
| def predict(self, x: RDD["VectorLike"]) -> RDD[float]: |
| ... |
| |
| def predict(self, x: Union["VectorLike", RDD["VectorLike"]]) -> Union[float, RDD[float]]: |
| """ |
| Predict the label of one or more examples. |
| |
| .. versionadded:: 1.1.0 |
| |
| Parameters |
| ---------- |
| x : :py:class:`pyspark.mllib.linalg.Vector` or :py:class:`pyspark.RDD` |
| Data point (feature vector), or an RDD of data points (feature |
| vectors). |
| |
| Notes |
| ----- |
| In Python, predict cannot currently be used within an RDD |
| transformation or action. |
| Call predict directly on the RDD instead. |
| """ |
| if isinstance(x, RDD): |
| return self.call("predict", x.map(_convert_to_vector)) |
| |
| else: |
| return self.call("predict", _convert_to_vector(x)) |
| |
| @since("1.1.0") |
| def numNodes(self) -> int: |
| """Get number of nodes in tree, including leaf nodes.""" |
| return self._java_model.numNodes() |
| |
| @since("1.1.0") |
| def depth(self) -> int: |
| """ |
| Get depth of tree (e.g. depth 0 means 1 leaf node, depth 1 |
| means 1 internal node + 2 leaf nodes). |
| """ |
| return self._java_model.depth() |
| |
| def __repr__(self) -> str: |
| """summary of model.""" |
| return self._java_model.toString() |
| |
| @since("1.2.0") |
| def toDebugString(self) -> str: |
| """full model.""" |
| return self._java_model.toDebugString() |
| |
| @classmethod |
| def _java_loader_class(cls) -> str: |
| return "org.apache.spark.mllib.tree.model.DecisionTreeModel" |
| |
| |
| class DecisionTree: |
| """ |
| Learning algorithm for a decision tree model for classification or |
| regression. |
| |
| .. versionadded:: 1.1.0 |
| """ |
| |
| @classmethod |
| def _train( |
| cls, |
| data: RDD[LabeledPoint], |
| type: str, |
| numClasses: int, |
| features: Dict[int, int], |
| impurity: str = "gini", |
| maxDepth: int = 5, |
| maxBins: int = 32, |
| minInstancesPerNode: int = 1, |
| minInfoGain: float = 0.0, |
| ) -> DecisionTreeModel: |
| first = data.first() |
| assert isinstance(first, LabeledPoint), "the data should be RDD of LabeledPoint" |
| model = callMLlibFunc( |
| "trainDecisionTreeModel", |
| data, |
| type, |
| numClasses, |
| features, |
| impurity, |
| maxDepth, |
| maxBins, |
| minInstancesPerNode, |
| minInfoGain, |
| ) |
| return DecisionTreeModel(model) |
| |
| @classmethod |
| def trainClassifier( |
| cls, |
| data: RDD[LabeledPoint], |
| numClasses: int, |
| categoricalFeaturesInfo: Dict[int, int], |
| impurity: str = "gini", |
| maxDepth: int = 5, |
| maxBins: int = 32, |
| minInstancesPerNode: int = 1, |
| minInfoGain: float = 0.0, |
| ) -> DecisionTreeModel: |
| """ |
| Train a decision tree model for classification. |
| |
| .. versionadded:: 1.1.0 |
| |
| Parameters |
| ---------- |
| data : :py:class:`pyspark.RDD` |
| Training data: RDD of LabeledPoint. Labels should take values |
| {0, 1, ..., numClasses-1}. |
| numClasses : int |
| Number of classes for classification. |
| categoricalFeaturesInfo : dict |
| Map storing arity of categorical features. An entry (n -> k) |
| indicates that feature n is categorical with k categories |
| indexed from 0: {0, 1, ..., k-1}. |
| impurity : str, optional |
| Criterion used for information gain calculation. |
| Supported values: "gini" or "entropy". |
| (default: "gini") |
| maxDepth : int, optional |
| Maximum depth of tree (e.g. depth 0 means 1 leaf node, depth 1 |
| means 1 internal node + 2 leaf nodes). |
| (default: 5) |
| maxBins : int, optional |
| Number of bins used for finding splits at each node. |
| (default: 32) |
| minInstancesPerNode : int, optional |
| Minimum number of instances required at child nodes to create |
| the parent split. |
| (default: 1) |
| minInfoGain : float, optional |
| Minimum info gain required to create a split. |
| (default: 0.0) |
| |
| Returns |
| ------- |
| :py:class:`DecisionTreeModel` |
| |
| Examples |
| -------- |
| >>> from numpy import array |
| >>> from pyspark.mllib.regression import LabeledPoint |
| >>> from pyspark.mllib.tree import DecisionTree |
| >>> |
| >>> data = [ |
| ... LabeledPoint(0.0, [0.0]), |
| ... LabeledPoint(1.0, [1.0]), |
| ... LabeledPoint(1.0, [2.0]), |
| ... LabeledPoint(1.0, [3.0]) |
| ... ] |
| >>> model = DecisionTree.trainClassifier(sc.parallelize(data), 2, {}) |
| >>> print(model) |
| DecisionTreeModel classifier of depth 1 with 3 nodes |
| |
| >>> print(model.toDebugString()) |
| DecisionTreeModel classifier of depth 1 with 3 nodes |
| If (feature 0 <= 0.5) |
| Predict: 0.0 |
| Else (feature 0 > 0.5) |
| Predict: 1.0 |
| >>> model.predict(array([1.0])) |
| 1.0 |
| >>> model.predict(array([0.0])) |
| 0.0 |
| >>> rdd = sc.parallelize([[1.0], [0.0]]) |
| >>> model.predict(rdd).collect() |
| [1.0, 0.0] |
| """ |
| return cls._train( |
| data, |
| "classification", |
| numClasses, |
| categoricalFeaturesInfo, |
| impurity, |
| maxDepth, |
| maxBins, |
| minInstancesPerNode, |
| minInfoGain, |
| ) |
| |
| @classmethod |
| @since("1.1.0") |
| def trainRegressor( |
| cls, |
| data: RDD[LabeledPoint], |
| categoricalFeaturesInfo: Dict[int, int], |
| impurity: str = "variance", |
| maxDepth: int = 5, |
| maxBins: int = 32, |
| minInstancesPerNode: int = 1, |
| minInfoGain: float = 0.0, |
| ) -> DecisionTreeModel: |
| """ |
| Train a decision tree model for regression. |
| |
| Parameters |
| ---------- |
| data : :py:class:`pyspark.RDD` |
| Training data: RDD of LabeledPoint. Labels are real numbers. |
| categoricalFeaturesInfo : dict |
| Map storing arity of categorical features. An entry (n -> k) |
| indicates that feature n is categorical with k categories |
| indexed from 0: {0, 1, ..., k-1}. |
| impurity : str, optional |
| Criterion used for information gain calculation. |
| The only supported value for regression is "variance". |
| (default: "variance") |
| maxDepth : int, optional |
| Maximum depth of tree (e.g. depth 0 means 1 leaf node, depth 1 |
| means 1 internal node + 2 leaf nodes). |
| (default: 5) |
| maxBins : int, optional |
| Number of bins used for finding splits at each node. |
| (default: 32) |
| minInstancesPerNode : int, optional |
| Minimum number of instances required at child nodes to create |
| the parent split. |
| (default: 1) |
| minInfoGain : float, optional |
| Minimum info gain required to create a split. |
| (default: 0.0) |
| |
| Returns |
| ------- |
| :py:class:`DecisionTreeModel` |
| |
| Examples |
| -------- |
| >>> from pyspark.mllib.regression import LabeledPoint |
| >>> from pyspark.mllib.tree import DecisionTree |
| >>> from pyspark.mllib.linalg import SparseVector |
| >>> |
| >>> sparse_data = [ |
| ... LabeledPoint(0.0, SparseVector(2, {0: 0.0})), |
| ... LabeledPoint(1.0, SparseVector(2, {1: 1.0})), |
| ... LabeledPoint(0.0, SparseVector(2, {0: 0.0})), |
| ... LabeledPoint(1.0, SparseVector(2, {1: 2.0})) |
| ... ] |
| >>> |
| >>> model = DecisionTree.trainRegressor(sc.parallelize(sparse_data), {}) |
| >>> model.predict(SparseVector(2, {1: 1.0})) |
| 1.0 |
| >>> model.predict(SparseVector(2, {1: 0.0})) |
| 0.0 |
| >>> rdd = sc.parallelize([[0.0, 1.0], [0.0, 0.0]]) |
| >>> model.predict(rdd).collect() |
| [1.0, 0.0] |
| """ |
| return cls._train( |
| data, |
| "regression", |
| 0, |
| categoricalFeaturesInfo, |
| impurity, |
| maxDepth, |
| maxBins, |
| minInstancesPerNode, |
| minInfoGain, |
| ) |
| |
| |
| @inherit_doc |
| class RandomForestModel(TreeEnsembleModel, JavaLoader["RandomForestModel"]): |
| """ |
| Represents a random forest model. |
| |
| .. versionadded:: 1.2.0 |
| """ |
| |
| @classmethod |
| def _java_loader_class(cls) -> str: |
| return "org.apache.spark.mllib.tree.model.RandomForestModel" |
| |
| |
| class RandomForest: |
| """ |
| Learning algorithm for a random forest model for classification or |
| regression. |
| |
| .. versionadded:: 1.2.0 |
| """ |
| |
| supportedFeatureSubsetStrategies: Tuple[str, ...] = ("auto", "all", "sqrt", "log2", "onethird") |
| |
| @classmethod |
| def _train( |
| cls, |
| data: RDD[LabeledPoint], |
| algo: str, |
| numClasses: int, |
| categoricalFeaturesInfo: Dict[int, int], |
| numTrees: int, |
| featureSubsetStrategy: str, |
| impurity: str, |
| maxDepth: int, |
| maxBins: int, |
| seed: Optional[int], |
| ) -> RandomForestModel: |
| first = data.first() |
| assert isinstance(first, LabeledPoint), "the data should be RDD of LabeledPoint" |
| if featureSubsetStrategy not in cls.supportedFeatureSubsetStrategies: |
| raise ValueError("unsupported featureSubsetStrategy: %s" % featureSubsetStrategy) |
| if seed is None: |
| seed = random.randint(0, 1 << 30) |
| model = callMLlibFunc( |
| "trainRandomForestModel", |
| data, |
| algo, |
| numClasses, |
| categoricalFeaturesInfo, |
| numTrees, |
| featureSubsetStrategy, |
| impurity, |
| maxDepth, |
| maxBins, |
| seed, |
| ) |
| return RandomForestModel(model) |
| |
| @classmethod |
| def trainClassifier( |
| cls, |
| data: RDD[LabeledPoint], |
| numClasses: int, |
| categoricalFeaturesInfo: Dict[int, int], |
| numTrees: int, |
| featureSubsetStrategy: str = "auto", |
| impurity: str = "gini", |
| maxDepth: int = 4, |
| maxBins: int = 32, |
| seed: Optional[int] = None, |
| ) -> RandomForestModel: |
| """ |
| Train a random forest model for binary or multiclass |
| classification. |
| |
| .. versionadded:: 1.2.0 |
| |
| Parameters |
| ---------- |
| data : :py:class:`pyspark.RDD` |
| Training dataset: RDD of LabeledPoint. Labels should take values |
| {0, 1, ..., numClasses-1}. |
| numClasses : int |
| Number of classes for classification. |
| categoricalFeaturesInfo : dict |
| Map storing arity of categorical features. An entry (n -> k) |
| indicates that feature n is categorical with k categories |
| indexed from 0: {0, 1, ..., k-1}. |
| numTrees : int |
| Number of trees in the random forest. |
| featureSubsetStrategy : str, optional |
| Number of features to consider for splits at each node. |
| Supported values: "auto", "all", "sqrt", "log2", "onethird". |
| If "auto" is set, this parameter is set based on numTrees: |
| if numTrees == 1, set to "all"; |
| if numTrees > 1 (forest) set to "sqrt". |
| (default: "auto") |
| impurity : str, optional |
| Criterion used for information gain calculation. |
| Supported values: "gini" or "entropy". |
| (default: "gini") |
| maxDepth : int, optional |
| Maximum depth of tree (e.g. depth 0 means 1 leaf node, depth 1 |
| means 1 internal node + 2 leaf nodes). |
| (default: 4) |
| maxBins : int, optional |
| Maximum number of bins used for splitting features. |
| (default: 32) |
| seed : int, Optional |
| Random seed for bootstrapping and choosing feature subsets. |
| Set as None to generate seed based on system time. |
| (default: None) |
| |
| Returns |
| ------- |
| :py:class:`RandomForestModel` |
| that can be used for prediction. |
| |
| Examples |
| -------- |
| >>> from pyspark.mllib.regression import LabeledPoint |
| >>> from pyspark.mllib.tree import RandomForest |
| >>> |
| >>> data = [ |
| ... LabeledPoint(0.0, [0.0]), |
| ... LabeledPoint(0.0, [1.0]), |
| ... LabeledPoint(1.0, [2.0]), |
| ... LabeledPoint(1.0, [3.0]) |
| ... ] |
| >>> model = RandomForest.trainClassifier(sc.parallelize(data), 2, {}, 3, seed=42) |
| >>> model.numTrees() |
| 3 |
| >>> model.totalNumNodes() |
| 7 |
| >>> print(model) |
| TreeEnsembleModel classifier with 3 trees |
| >>> print(model.toDebugString()) |
| TreeEnsembleModel classifier with 3 trees |
| Tree 0: |
| Predict: 1.0 |
| Tree 1: |
| If (feature 0 <= 1.5) |
| Predict: 0.0 |
| Else (feature 0 > 1.5) |
| Predict: 1.0 |
| Tree 2: |
| If (feature 0 <= 1.5) |
| Predict: 0.0 |
| Else (feature 0 > 1.5) |
| Predict: 1.0 |
| >>> model.predict([2.0]) |
| 1.0 |
| >>> model.predict([0.0]) |
| 0.0 |
| >>> rdd = sc.parallelize([[3.0], [1.0]]) |
| >>> model.predict(rdd).collect() |
| [1.0, 0.0] |
| """ |
| return cls._train( |
| data, |
| "classification", |
| numClasses, |
| categoricalFeaturesInfo, |
| numTrees, |
| featureSubsetStrategy, |
| impurity, |
| maxDepth, |
| maxBins, |
| seed, |
| ) |
| |
| @classmethod |
| def trainRegressor( |
| cls, |
| data: RDD[LabeledPoint], |
| categoricalFeaturesInfo: Dict[int, int], |
| numTrees: int, |
| featureSubsetStrategy: str = "auto", |
| impurity: str = "variance", |
| maxDepth: int = 4, |
| maxBins: int = 32, |
| seed: Optional[int] = None, |
| ) -> RandomForestModel: |
| """ |
| Train a random forest model for regression. |
| |
| .. versionadded:: 1.2.0 |
| |
| Parameters |
| ---------- |
| data : :py:class:`pyspark.RDD` |
| Training dataset: RDD of LabeledPoint. Labels are real numbers. |
| categoricalFeaturesInfo : dict |
| Map storing arity of categorical features. An entry (n -> k) |
| indicates that feature n is categorical with k categories |
| indexed from 0: {0, 1, ..., k-1}. |
| numTrees : int |
| Number of trees in the random forest. |
| featureSubsetStrategy : str, optional |
| Number of features to consider for splits at each node. |
| Supported values: "auto", "all", "sqrt", "log2", "onethird". |
| If "auto" is set, this parameter is set based on numTrees: |
| |
| - if numTrees == 1, set to "all"; |
| - if numTrees > 1 (forest) set to "onethird" for regression. |
| |
| (default: "auto") |
| impurity : str, optional |
| Criterion used for information gain calculation. |
| The only supported value for regression is "variance". |
| (default: "variance") |
| maxDepth : int, optional |
| Maximum depth of tree (e.g. depth 0 means 1 leaf node, depth 1 |
| means 1 internal node + 2 leaf nodes). |
| (default: 4) |
| maxBins : int, optional |
| Maximum number of bins used for splitting features. |
| (default: 32) |
| seed : int, optional |
| Random seed for bootstrapping and choosing feature subsets. |
| Set as None to generate seed based on system time. |
| (default: None) |
| |
| Returns |
| ------- |
| :py:class:`RandomForestModel` |
| that can be used for prediction. |
| |
| Examples |
| -------- |
| >>> from pyspark.mllib.regression import LabeledPoint |
| >>> from pyspark.mllib.tree import RandomForest |
| >>> from pyspark.mllib.linalg import SparseVector |
| >>> |
| >>> sparse_data = [ |
| ... LabeledPoint(0.0, SparseVector(2, {0: 1.0})), |
| ... LabeledPoint(1.0, SparseVector(2, {1: 1.0})), |
| ... LabeledPoint(0.0, SparseVector(2, {0: 1.0})), |
| ... LabeledPoint(1.0, SparseVector(2, {1: 2.0})) |
| ... ] |
| >>> |
| >>> model = RandomForest.trainRegressor(sc.parallelize(sparse_data), {}, 2, seed=42) |
| >>> model.numTrees() |
| 2 |
| >>> model.totalNumNodes() |
| 4 |
| >>> model.predict(SparseVector(2, {1: 1.0})) |
| 1.0 |
| >>> model.predict(SparseVector(2, {0: 1.0})) |
| 0.5 |
| >>> rdd = sc.parallelize([[0.0, 1.0], [1.0, 0.0]]) |
| >>> model.predict(rdd).collect() |
| [1.0, 0.5] |
| """ |
| return cls._train( |
| data, |
| "regression", |
| 0, |
| categoricalFeaturesInfo, |
| numTrees, |
| featureSubsetStrategy, |
| impurity, |
| maxDepth, |
| maxBins, |
| seed, |
| ) |
| |
| |
| @inherit_doc |
| class GradientBoostedTreesModel(TreeEnsembleModel, JavaLoader["GradientBoostedTreesModel"]): |
| """ |
| Represents a gradient-boosted tree model. |
| |
| .. versionadded:: 1.3.0 |
| """ |
| |
| @classmethod |
| def _java_loader_class(cls) -> str: |
| return "org.apache.spark.mllib.tree.model.GradientBoostedTreesModel" |
| |
| |
| class GradientBoostedTrees: |
| """ |
| Learning algorithm for a gradient boosted trees model for |
| classification or regression. |
| |
| .. versionadded:: 1.3.0 |
| """ |
| |
| @classmethod |
| def _train( |
| cls, |
| data: RDD[LabeledPoint], |
| algo: str, |
| categoricalFeaturesInfo: Dict[int, int], |
| loss: str, |
| numIterations: int, |
| learningRate: float, |
| maxDepth: int, |
| maxBins: int, |
| ) -> GradientBoostedTreesModel: |
| first = data.first() |
| assert isinstance(first, LabeledPoint), "the data should be RDD of LabeledPoint" |
| model = callMLlibFunc( |
| "trainGradientBoostedTreesModel", |
| data, |
| algo, |
| categoricalFeaturesInfo, |
| loss, |
| numIterations, |
| learningRate, |
| maxDepth, |
| maxBins, |
| ) |
| return GradientBoostedTreesModel(model) |
| |
| @classmethod |
| def trainClassifier( |
| cls, |
| data: RDD[LabeledPoint], |
| categoricalFeaturesInfo: Dict[int, int], |
| loss: str = "logLoss", |
| numIterations: int = 100, |
| learningRate: float = 0.1, |
| maxDepth: int = 3, |
| maxBins: int = 32, |
| ) -> GradientBoostedTreesModel: |
| """ |
| Train a gradient-boosted trees model for classification. |
| |
| .. versionadded:: 1.3.0 |
| |
| Parameters |
| ---------- |
| data : :py:class:`pyspark.RDD` |
| Training dataset: RDD of LabeledPoint. Labels should take values |
| {0, 1}. |
| categoricalFeaturesInfo : dict |
| Map storing arity of categorical features. An entry (n -> k) |
| indicates that feature n is categorical with k categories |
| indexed from 0: {0, 1, ..., k-1}. |
| loss : str, optional |
| Loss function used for minimization during gradient boosting. |
| Supported values: "logLoss", "leastSquaresError", |
| "leastAbsoluteError". |
| (default: "logLoss") |
| numIterations : int, optional |
| Number of iterations of boosting. |
| (default: 100) |
| learningRate : float, optional |
| Learning rate for shrinking the contribution of each estimator. |
| The learning rate should be between in the interval (0, 1]. |
| (default: 0.1) |
| maxDepth : int, optional |
| Maximum depth of tree (e.g. depth 0 means 1 leaf node, depth 1 |
| means 1 internal node + 2 leaf nodes). |
| (default: 3) |
| maxBins : int, optional |
| Maximum number of bins used for splitting features. DecisionTree |
| requires maxBins >= max categories. |
| (default: 32) |
| |
| Returns |
| ------- |
| :py:class:`GradientBoostedTreesModel` |
| that can be used for prediction. |
| |
| Examples |
| -------- |
| >>> from pyspark.mllib.regression import LabeledPoint |
| >>> from pyspark.mllib.tree import GradientBoostedTrees |
| >>> |
| >>> data = [ |
| ... LabeledPoint(0.0, [0.0]), |
| ... LabeledPoint(0.0, [1.0]), |
| ... LabeledPoint(1.0, [2.0]), |
| ... LabeledPoint(1.0, [3.0]) |
| ... ] |
| >>> |
| >>> model = GradientBoostedTrees.trainClassifier(sc.parallelize(data), {}, numIterations=10) |
| >>> model.numTrees() |
| 10 |
| >>> model.totalNumNodes() |
| 30 |
| >>> print(model) # it already has newline |
| TreeEnsembleModel classifier with 10 trees |
| >>> model.predict([2.0]) |
| 1.0 |
| >>> model.predict([0.0]) |
| 0.0 |
| >>> rdd = sc.parallelize([[2.0], [0.0]]) |
| >>> model.predict(rdd).collect() |
| [1.0, 0.0] |
| """ |
| return cls._train( |
| data, |
| "classification", |
| categoricalFeaturesInfo, |
| loss, |
| numIterations, |
| learningRate, |
| maxDepth, |
| maxBins, |
| ) |
| |
| @classmethod |
| def trainRegressor( |
| cls, |
| data: RDD[LabeledPoint], |
| categoricalFeaturesInfo: Dict[int, int], |
| loss: str = "leastSquaresError", |
| numIterations: int = 100, |
| learningRate: float = 0.1, |
| maxDepth: int = 3, |
| maxBins: int = 32, |
| ) -> GradientBoostedTreesModel: |
| """ |
| Train a gradient-boosted trees model for regression. |
| |
| .. versionadded:: 1.3.0 |
| |
| Parameters |
| ---------- |
| data : |
| Training dataset: RDD of LabeledPoint. Labels are real numbers. |
| categoricalFeaturesInfo : dict |
| Map storing arity of categorical features. An entry (n -> k) |
| indicates that feature n is categorical with k categories |
| indexed from 0: {0, 1, ..., k-1}. |
| loss : str, optional |
| Loss function used for minimization during gradient boosting. |
| Supported values: "logLoss", "leastSquaresError", |
| "leastAbsoluteError". |
| (default: "leastSquaresError") |
| numIterations : int, optional |
| Number of iterations of boosting. |
| (default: 100) |
| learningRate : float, optional |
| Learning rate for shrinking the contribution of each estimator. |
| The learning rate should be between in the interval (0, 1]. |
| (default: 0.1) |
| maxDepth : int, optional |
| Maximum depth of tree (e.g. depth 0 means 1 leaf node, depth 1 |
| means 1 internal node + 2 leaf nodes). |
| (default: 3) |
| maxBins : int, optional |
| Maximum number of bins used for splitting features. DecisionTree |
| requires maxBins >= max categories. |
| (default: 32) |
| |
| Returns |
| ------- |
| :py:class:`GradientBoostedTreesModel` |
| that can be used for prediction. |
| |
| Examples |
| -------- |
| >>> from pyspark.mllib.regression import LabeledPoint |
| >>> from pyspark.mllib.tree import GradientBoostedTrees |
| >>> from pyspark.mllib.linalg import SparseVector |
| >>> |
| >>> sparse_data = [ |
| ... LabeledPoint(0.0, SparseVector(2, {0: 1.0})), |
| ... LabeledPoint(1.0, SparseVector(2, {1: 1.0})), |
| ... LabeledPoint(0.0, SparseVector(2, {0: 1.0})), |
| ... LabeledPoint(1.0, SparseVector(2, {1: 2.0})) |
| ... ] |
| >>> |
| >>> data = sc.parallelize(sparse_data) |
| >>> model = GradientBoostedTrees.trainRegressor(data, {}, numIterations=10) |
| >>> model.numTrees() |
| 10 |
| >>> model.totalNumNodes() |
| 12 |
| >>> model.predict(SparseVector(2, {1: 1.0})) |
| 1.0 |
| >>> model.predict(SparseVector(2, {0: 1.0})) |
| 0.0 |
| >>> rdd = sc.parallelize([[0.0, 1.0], [1.0, 0.0]]) |
| >>> model.predict(rdd).collect() |
| [1.0, 0.0] |
| """ |
| return cls._train( |
| data, |
| "regression", |
| categoricalFeaturesInfo, |
| loss, |
| numIterations, |
| learningRate, |
| maxDepth, |
| maxBins, |
| ) |
| |
| |
| def _test() -> None: |
| import doctest |
| |
| globs = globals().copy() |
| from pyspark.sql import SparkSession |
| |
| spark = SparkSession.builder.master("local[4]").appName("mllib.tree tests").getOrCreate() |
| globs["sc"] = spark.sparkContext |
| (failure_count, test_count) = doctest.testmod( |
| globs=globs, optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE |
| ) |
| spark.stop() |
| if failure_count: |
| sys.exit(-1) |
| |
| |
| if __name__ == "__main__": |
| _test() |
