From f3996fec6ac16e7f4c223cfc716a9e0d0409b23f Mon Sep 17 00:00:00 2001 From: Y-oHr-N Date: Thu, 20 Feb 2020 16:38:22 +0900 Subject: [PATCH] Fix _VotingBooster --- optgbm/sklearn.py | 22 ++++++---------------- 1 file changed, 6 insertions(+), 16 deletions(-) diff --git a/optgbm/sklearn.py b/optgbm/sklearn.py index afb0e0b..e6093c5 100644 --- a/optgbm/sklearn.py +++ b/optgbm/sklearn.py @@ -249,25 +249,15 @@ def from_representations( return cls(boosters, weights=weights) - def predict( - self, X: TWO_DIM_ARRAYLIKE_TYPE, **kwargs: Any - ) -> TWO_DIM_ARRAYLIKE_TYPE: - results = [] - - for b in self.boosters: - result = b.predict(X, **kwargs) - - results.append(result) - - return np.average(results, axis=0, weights=self.weights) - def feature_importance(self, **kwargs: Any) -> np.ndarray: - results = [] + results = [b.feature_importance(**kwargs) for b in self.boosters] - for b in self.boosters: - result = b.feature_importance(**kwargs) + return np.average(results, axis=0, weights=self.weights) - results.append(result) + def predict( + self, X: TWO_DIM_ARRAYLIKE_TYPE, **kwargs: Any + ) -> TWO_DIM_ARRAYLIKE_TYPE: + results = [b.predict(X, **kwargs) for b in self.boosters] return np.average(results, axis=0, weights=self.weights)