More general predict proba. (#6817)

* Use `output_margin` for `softmax`.
* Add test for dask binary cls.

Co-authored-by: Philip Hyunsu Cho <chohyu01@cs.washington.edu>

python-package/xgboost/dask.py

@@ -21,6 +21,7 @@ from contextlib import contextmanager
 from collections import defaultdict
 from collections.abc import Sequence
 from threading import Thread
+from functools import partial, update_wrapper
 from typing import TYPE_CHECKING, List, Tuple, Callable, Optional, Any, Union, Dict, Set
 from typing import Awaitable, Generator, TypeVar
 
@@ -968,7 +969,7 @@ def _can_output_df(is_df: bool, output_shape: Tuple) -> bool:
     return is_df and len(output_shape) <= 2
 
 
-async def _direct_predict_impl(
+async def _direct_predict_impl(  # pylint: disable=too-many-branches
     mapped_predict: Callable,
     booster: "distributed.Future",
     data: _DaskCollection,
@@ -1023,6 +1024,14 @@ async def _direct_predict_impl(
                 new_axis = list(range(len(output_shape) - 2))
             else:
                 new_axis = [i + 2 for i in range(len(output_shape) - 2)]
+        if len(output_shape) == 2:
+            # Somehow dask fail to infer output shape change for 2-dim prediction, and
+            #  `chunks = (None, output_shape[1])` doesn't work due to None is not
+            #  supported in map_blocks.
+            chunks = list(data.chunks)
+            chunks[1] = (output_shape[1], )
+        else:
+            chunks = None
         predictions = da.map_blocks(
             mapped_predict,
             booster,
@@ -1030,6 +1039,8 @@ async def _direct_predict_impl(
             False,
             columns,
             base_margin_array,
+
+            chunks=chunks,
             drop_axis=drop_axis,
             new_axis=new_axis,
             dtype=numpy.float32,
@@ -1777,20 +1788,20 @@ class DaskXGBClassifier(DaskScikitLearnBase, XGBClassifierBase):
         self,
         X: _DaskCollection,
         validate_features: bool,
-        output_margin: bool,
         base_margin: Optional[_DaskCollection],
         iteration_range: Optional[Tuple[int, int]],
     ) -> _DaskCollection:
-        if iteration_range is None:
-            iteration_range = (0, 0)
         predts = await super()._predict_async(
             data=X,
-            output_margin=output_margin,
+            output_margin=self.objective == "multi:softmax",
             validate_features=validate_features,
             base_margin=base_margin,
             iteration_range=iteration_range,
         )
-        return _cls_predict_proba(self.objective, predts, da.vstack)
+        vstack = update_wrapper(
+            partial(da.vstack, allow_unknown_chunksizes=True), da.vstack
+        )
+        return _cls_predict_proba(getattr(self, "n_classes_", None), predts, vstack)
 
     # pylint: disable=missing-function-docstring
     def predict_proba(
@@ -1798,7 +1809,6 @@ class DaskXGBClassifier(DaskScikitLearnBase, XGBClassifierBase):
         X: _DaskCollection,
         ntree_limit: Optional[int] = None,
         validate_features: bool = True,
-        output_margin: bool = False,
         base_margin: Optional[_DaskCollection] = None,
         iteration_range: Optional[Tuple[int, int]] = None,
     ) -> Any:
@@ -1809,7 +1819,6 @@ class DaskXGBClassifier(DaskScikitLearnBase, XGBClassifierBase):
             self._predict_proba_async,
             X=X,
             validate_features=validate_features,
-            output_margin=output_margin,
             base_margin=base_margin,
             iteration_range=iteration_range,
         )

python-package/xgboost/sklearn.py

@@ -5,7 +5,7 @@ import copy
 import warnings
 import json
 import os
-from typing import Union, Optional, List, Dict, Callable, Tuple, Any
+from typing import Union, Optional, List, Dict, Callable, Tuple, Any, TypeVar
 import numpy as np
 from .core import Booster, DMatrix, XGBoostError
 from .core import _deprecate_positional_args, _convert_ntree_limit
@@ -534,6 +534,8 @@ class XGBModel(XGBModelBase):
         self.get_booster().load_model(fname)
         meta = self.get_booster().attr('scikit_learn')
         if meta is None:
+            # FIXME(jiaming): This doesn't have to be a problem as most of the needed
+            # information like num_class and objective is in Learner class.
             warnings.warn(
                 'Loading a native XGBoost model with Scikit-Learn interface.'
             )
@@ -545,6 +547,8 @@ class XGBModel(XGBModelBase):
                 self._le = XGBoostLabelEncoder()
                 self._le.from_json(v)
                 continue
+            # FIXME(jiaming): This can be removed once label encoder is gone since we can
+            # generate it from `np.arange(self.n_classes_)`
             if k == 'classes_':
                 self.classes_ = np.array(v)
                 continue
@@ -1000,17 +1004,14 @@ class XGBModel(XGBModelBase):
         return np.array(json.loads(b.get_dump(dump_format='json')[0])['bias'])
 
 
-def _cls_predict_proba(
-    objective: Union[str, Callable], prediction: Any, vstack: Callable
-) -> Any:
-    if objective == 'multi:softmax':
-        raise ValueError('multi:softmax objective does not support predict_proba,'
-                         ' use `multi:softprob` or `binary:logistic` instead.')
-    if objective == 'multi:softprob' or callable(objective):
-        # Return prediction directly if if objective is defined by user since we don't
-        # know how to perform the transformation
+PredtT = TypeVar("PredtT")
+
+
+def _cls_predict_proba(n_classes: int, prediction: PredtT, vstack: Callable) -> PredtT:
+    assert len(prediction.shape) <= 2
+    if len(prediction.shape) == 2 and prediction.shape[1] == n_classes:
         return prediction
-    # Lastly the binary logistic function
+    # binary logistic function
     classone_probs = prediction
     classzero_probs = 1.0 - classone_probs
     return vstack((classzero_probs, classone_probs)).transpose()
@@ -1194,8 +1195,10 @@ class XGBClassifier(XGBModel, XGBClassifierBase):
             return class_probs
 
         if len(class_probs.shape) > 1:
+            # turns softprob into softmax
             column_indexes = np.argmax(class_probs, axis=1)
         else:
+            # turns soft logit into class label
             column_indexes = np.repeat(0, class_probs.shape[0])
             column_indexes[class_probs > 0.5] = 1
 
@@ -1238,15 +1241,23 @@ class XGBClassifier(XGBModel, XGBClassifierBase):
             a numpy array of shape array-like of shape (n_samples, n_classes) with the
             probability of each data example being of a given class.
         """
+        # custom obj:      Do nothing as we don't know what to do.
+        # softprob:        Do nothing, output is proba.
+        # softmax:         Use output margin to remove the argmax in PredTransform.
+        # binary:logistic: Expand the prob vector into 2-class matrix after predict.
+        # binary:logitraw: Unsupported by predict_proba()
         class_probs = super().predict(
             X=X,
-            output_margin=False,
+            output_margin=self.objective == "multi:softmax",
             ntree_limit=ntree_limit,
             validate_features=validate_features,
             base_margin=base_margin,
             iteration_range=iteration_range
         )
-        return _cls_predict_proba(self.objective, class_probs, np.vstack)
+        # If model is loaded from a raw booster there's no `n_classes_`
+        return _cls_predict_proba(
+            getattr(self, "n_classes_", None), class_probs, np.vstack
+        )
 
     def evals_result(self):
         """Return the evaluation results.