Hendrik/xgboost
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

ENH: Add n_classes_ to fitted classifier.

Browse Source
This commit is contained in:
Skipper Seabold 2015-04-27 10:46:30 -05:00
parent 8ac89b290e
commit c1a24c0fb1
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
wrapper/xgboost.py
View File

@ -872,11 +872,12 @@ class XGBClassifier(XGBModel, XGBClassifier):
def fit(self, X, y, sample_weight=None): def fit(self, X, y, sample_weight=None):
y_values = list(np.unique(y)) y_values = list(np.unique(y))
if len(y_values) > 2: self.n_classes_ = len(y_values)
if self.n_classes_ > 2:
# Switch to using a multiclass objective in the underlying XGB instance # Switch to using a multiclass objective in the underlying XGB instance
self.objective = "multi:softprob" self.objective = "multi:softprob"
xgb_options = self.get_xgb_params() xgb_options = self.get_xgb_params()
xgb_options['num_class'] = len(y_values) xgb_options['num_class'] = self.n_classes_
else: else:
xgb_options = self.get_xgb_params() xgb_options = self.get_xgb_params()