decision_tree.py 文件源码

def fit(self, X, y, sample_weight=None):
        from sklearn.tree import DecisionTreeRegressor

        self.max_features = float(self.max_features)
        if self.max_depth == "None":
            self.max_depth = None
        else:
            num_features = X.shape[1]
            max_depth = max(1, int(np.round(self.max_depth * num_features, 0)))
        self.min_samples_split = int(self.min_samples_split)
        self.min_samples_leaf = int(self.min_samples_leaf)
        if self.max_leaf_nodes == "None":
            self.max_leaf_nodes = None
        else:
            self.max_leaf_nodes = int(self.max_leaf_nodes)
        self.min_weight_fraction_leaf = float(self.min_weight_fraction_leaf)

        self.estimator = DecisionTreeRegressor(
            criterion=self.criterion,
            max_depth=max_depth,
            min_samples_split=self.min_samples_split,
            min_samples_leaf=self.min_samples_leaf,
            max_leaf_nodes=self.max_leaf_nodes,
            random_state=self.random_state)
        self.estimator.fit(X, y, sample_weight=sample_weight)
        return self