Repository: onnx/sklearn-onnx

Issue:
Implement converter for BayesianGaussianMixture



Discussion:

Resolving patch:
diff --git a/skl2onnx/operator_converters/gaussian_mixture.py b/skl2onnx/operator_converters/gaussian_mixture.py
--- a/skl2onnx/operator_converters/gaussian_mixture.py
+++ b/skl2onnx/operator_converters/gaussian_mixture.py
@@ -5,13 +5,15 @@
 # --------------------------------------------------------------------------
 
 import numpy as np
-from sklearn.mixture import BayesianGaussianMixture
+from scipy.special import digamma
+from sklearn.mixture import BayesianGaussianMixture, GaussianMixture
 try:
     from sklearn.mixture._gaussian_mixture import _compute_log_det_cholesky
 except ImportError:
     # scikit-learn < 0.22
     from sklearn.mixture.gaussian_mixture import _compute_log_det_cholesky
 from ..common._registration import register_converter
+from ..common.data_types import guess_numpy_type
 from ..algebra.onnx_ops import (
     OnnxAdd, OnnxSub, OnnxMul, OnnxGemm, OnnxReduceSumSquare,
     OnnxReduceLogSumExp, OnnxExp, OnnxArgMax, OnnxConcat,
@@ -20,42 +22,23 @@
 from ..proto import onnx_proto
 
 
-def convert_sklearn_gaussian_mixture(scope, operator, container):
+def _estimate_log_gaussian_prob(X, means, precisions_chol,
+                                covariance_type, dtype, op_version,
+                                combined_reducesum):
     """
-    Converter for *GaussianMixture*,
-    *BayesianGaussianMixture*.
-    Parameters which change the prediction function:
-
-    * *covariance_type*
+    Converts the same function into ONNX.
+    Returns log probabilities.
     """
-    X = operator.inputs[0]
-    out = operator.outputs
-    op = operator.raw_operator
-    n_features = X.type.shape[1]
-    n_components = op.means_.shape[0]
-    opv = container.target_opset
-    options = container.get_options(op, dict(score_samples=None))
-    add_score = options.get('score_samples', False)
-    combined_reducesum = not container.is_allowed(
-        {'ReduceLogSumExp', 'ReduceSumSquare'})
-    if add_score and len(out) != 3:
-        raise RuntimeError("3 outputs are expected.")
-    if isinstance(op, BayesianGaussianMixture):
-        raise NotImplementedError(
-            "Converter for BayesianGaussianMixture is not implemented.")
-
-    # All comments come from scikit-learn code and tells
-    # which functions is being onnxified.
-    # def _estimate_weighted_log_prob(self, X):
-    # self._estimate_log_prob(X) + self._estimate_log_weights()
-    log_weights = np.log(op.weights_)  # self._estimate_log_weights()
+    n_components = means.shape[0]
+    n_features = means.shape[1]
+    opv = op_version
 
     # self._estimate_log_prob(X)
     log_det = _compute_log_det_cholesky(
-        op.precisions_cholesky_, op.covariance_type, n_features).astype(
-            container.dtype)
+        precisions_chol, covariance_type, n_features).astype(
+            dtype)
 
-    if op.covariance_type == 'full':
+    if covariance_type == 'full':
         # shape(op.means_) = (n_components, n_features)
         # shape(op.precisions_cholesky_) =
         #   (n_components, n_features, n_features)
@@ -67,10 +50,10 @@ def convert_sklearn_gaussian_mixture(scope, operator, container):
 
         ys = []
         for c in range(n_components):
-            prec_chol = op.precisions_cholesky_[c, :, :]
-            cst = - np.dot(op.means_[c, :], prec_chol)
-            y = OnnxGemm(X, prec_chol.astype(container.dtype),
-                         cst.astype(container.dtype), alpha=1.,
+            prec_chol = precisions_chol[c, :, :]
+            cst = - np.dot(means[c, :], prec_chol)
+            y = OnnxGemm(X, prec_chol.astype(dtype),
+                         cst.astype(dtype), alpha=1.,
                          beta=1., op_version=opv)
             if combined_reducesum:
                 y2s = OnnxReduceSum(OnnxMul(y, y, op_version=opv),
@@ -80,7 +63,7 @@ def convert_sklearn_gaussian_mixture(scope, operator, container):
             ys.append(y2s)
         log_prob = OnnxConcat(*ys, axis=1, op_version=opv)
 
-    elif op.covariance_type == 'tied':
+    elif covariance_type == 'tied':
         # shape(op.means_) = (n_components, n_features)
         # shape(op.precisions_cholesky_) =
         #   (n_features, n_features)
@@ -90,12 +73,11 @@ def convert_sklearn_gaussian_mixture(scope, operator, container):
         #     y = np.dot(X, precisions_chol) - np.dot(mu, precisions_chol)
         #     log_prob[:, k] = np.sum(np.square(y), axis=1)
 
-        precisions_chol = op.precisions_cholesky_
         ys = []
         for f in range(n_components):
-            cst = - np.dot(op.means_[f, :], precisions_chol)
-            y = OnnxGemm(X, precisions_chol.astype(container.dtype),
-                         cst.astype(container.dtype),
+            cst = - np.dot(means[f, :], precisions_chol)
+            y = OnnxGemm(X, precisions_chol.astype(dtype),
+                         cst.astype(dtype),
                          alpha=1., beta=1., op_version=opv)
             if combined_reducesum:
                 y2s = OnnxReduceSum(OnnxMul(y, y, op_version=opv),
@@ -105,7 +87,7 @@ def convert_sklearn_gaussian_mixture(scope, operator, container):
             ys.append(y2s)
         log_prob = OnnxConcat(*ys, axis=1, op_version=opv)
 
-    elif op.covariance_type == 'diag':
+    elif covariance_type == 'diag':
         # shape(op.means_) = (n_components, n_features)
         # shape(op.precisions_cholesky_) =
         #   (n_components, n_features)
@@ -115,20 +97,20 @@ def convert_sklearn_gaussian_mixture(scope, operator, container):
         #             2. * np.dot(X, (means * precisions).T) +
         #             np.dot(X ** 2, precisions.T))
 
-        precisions = op.precisions_cholesky_ ** 2
-        mp = np.sum((op.means_ ** 2 * precisions), 1)
-        zeros = np.zeros((n_components, ), dtype=container.dtype)
+        precisions = (precisions_chol ** 2).astype(dtype)
+        mp = np.sum((means ** 2 * precisions), 1).astype(dtype)
+        zeros = np.zeros((n_components, ), dtype=dtype)
         xmp = OnnxGemm(
-            X, (op.means_ * precisions).T.astype(container.dtype),
+            X, (means * precisions).T.astype(dtype),
             zeros, alpha=-2., beta=0., op_version=opv)
         term = OnnxGemm(OnnxMul(X, X, op_version=opv),
-                        precisions.T.astype(container.dtype),
+                        precisions.T.astype(dtype),
                         zeros, alpha=1., beta=0., op_version=opv)
         log_prob = OnnxAdd(
-            OnnxAdd(mp.astype(container.dtype), xmp, op_version=opv),
+            OnnxAdd(mp.astype(dtype), xmp, op_version=opv),
             term, op_version=opv)
 
-    elif op.covariance_type == 'spherical':
+    elif covariance_type == 'spherical':
         # shape(op.means_) = (n_components, n_features)
         # shape(op.precisions_cholesky_) = (n_components, )
 
@@ -137,40 +119,96 @@ def convert_sklearn_gaussian_mixture(scope, operator, container):
         #             2 * np.dot(X, means.T * precisions) +
         #             np.outer(row_norms(X, squared=True), precisions))
 
-        zeros = np.zeros((n_components, ))
-        precisions = op.precisions_cholesky_ ** 2
+        zeros = np.zeros((n_components, ), dtype=dtype)
+        precisions = (precisions_chol ** 2).astype(dtype)
         if combined_reducesum:
             normX = OnnxReduceSum(OnnxMul(X, X, op_version=opv),
                                   axes=[1], op_version=opv)
         else:
             normX = OnnxReduceSumSquare(X, axes=[1], op_version=opv)
         outer = OnnxGemm(
-            normX, precisions[np.newaxis, :].astype(container.dtype),
-            zeros.astype(container.dtype), alpha=1., beta=1., op_version=opv)
+            normX, precisions[np.newaxis, :].astype(dtype),
+            zeros.astype(dtype), alpha=1., beta=1., op_version=opv)
         xmp = OnnxGemm(
-            X, (op.means_.T * precisions).astype(container.dtype),
-            zeros.astype(container.dtype), alpha=-2.,
-            beta=0., op_version=opv)
-        mp = np.sum(op.means_ ** 2, 1) * precisions
-        log_prob = OnnxAdd(
-            mp.astype(container.dtype),
-            OnnxAdd(xmp, outer, op_version=opv),
-            op_version=opv)
+            X, (means.T * precisions).astype(dtype),
+            zeros, alpha=-2., beta=0., op_version=opv)
+        mp = (np.sum(means ** 2, 1) * precisions).astype(dtype)
+        log_prob = OnnxAdd(mp, OnnxAdd(xmp, outer, op_version=opv),
+                           op_version=opv)
     else:
         raise RuntimeError("Unknown op.covariance_type='{}'. Upgrade "
                            "to a more recent version of skearn-onnx "
-                           "or raise an issue.".format(op.covariance_type))
-
+                           "or raise an issue.".format(covariance_type))
     # -.5 * (cst + log_prob) + log_det
-    cst = np.array([n_features * np.log(2 * np.pi)])
-    add = OnnxAdd(cst.astype(container.dtype), log_prob, op_version=opv)
-    mul = OnnxMul(add, np.array([-0.5], dtype=container.dtype),
+    cst = np.array([n_features * np.log(2 * np.pi)]).astype(dtype)
+    add = OnnxAdd(cst, log_prob, op_version=opv)
+    mul = OnnxMul(add, np.array([-0.5], dtype=dtype),
                   op_version=opv)
     if isinstance(log_det, float):
-        log_det = np.array([log_det], dtype=container.dtype)
-    weighted_log_prob = OnnxAdd(
-        OnnxAdd(mul, log_det, op_version=opv),
-        log_weights.astype(container.dtype), op_version=opv)
+        log_det = np.array([log_det], dtype=dtype)
+
+    return OnnxAdd(mul, log_det.astype(dtype), op_version=opv)
+
+
+def convert_sklearn_gaussian_mixture(scope, operator, container):
+    """
+    Converter for *GaussianMixture*,
+    *BayesianGaussianMixture*.
+    """
+    X = operator.inputs[0]
+    dtype = guess_numpy_type(X.type)
+    if dtype != np.float64:
+        dtype = np.float32
+    out = operator.outputs
+    op = operator.raw_operator
+    n_components = op.means_.shape[0]
+    opv = container.target_opset
+    options = container.get_options(op, dict(score_samples=None))
+    add_score = options.get('score_samples', False)
+    combined_reducesum = not container.is_allowed(
+        {'ReduceLogSumExp', 'ReduceSumSquare'})
+    if add_score and len(out) != 3:
+        raise RuntimeError("3 outputs are expected.")
+
+    if X.type is not None:
+        if X.type.shape[1] != op.means_.shape[1]:
+            raise RuntimeError(
+                "Dimension mismath between expected number of features {} "
+                "and ONNX graphs expectations {}.".format(
+                    op.means_.shape[1], X.type.shape[1]))
+    n_features = op.means_.shape[1]
+
+    # All comments come from scikit-learn code and tells
+    # which functions is being onnxified.
+    # def _estimate_weighted_log_prob(self, X):
+    log_weights = op._estimate_log_weights().astype(dtype)
+
+    log_gauss = _estimate_log_gaussian_prob(
+        X, op.means_, op.precisions_cholesky_, op.covariance_type,
+        dtype, opv, combined_reducesum)
+
+    if isinstance(op, BayesianGaussianMixture):
+        # log_gauss = (_estimate_log_gaussian_prob(
+        #   X, self.means_, self.precisions_cholesky_, self.covariance_type) -
+        #   .5 * n_features * np.log(self.degrees_of_freedom_))
+
+        log_lambda = n_features * np.log(2.) + np.sum(digamma(
+            .5 * (op.degrees_of_freedom_ -
+                  np.arange(0, n_features)[:, np.newaxis])), 0)
+        cst_log_lambda = .5 * (log_lambda - n_features / op.mean_precision_)
+        cst = cst_log_lambda - .5 * n_features * np.log(op.degrees_of_freedom_)
+        if isinstance(cst, np.ndarray):
+            cst_array = cst.astype(dtype)
+        else:
+            cst_array = np.array([cst], dtype=dtype)
+        log_gauss = OnnxAdd(log_gauss, cst_array, op_version=opv)
+    elif not isinstance(op, GaussianMixture):
+        raise RuntimeError(
+            "The converter does not support type {}.".format(
+                type(op)))
+
+    # self._estimate_log_prob(X) + self._estimate_log_weights()
+    weighted_log_prob = OnnxAdd(log_gauss, log_weights, op_version=opv)
 
     # labels
     if container.is_allowed('ArgMax'):
@@ -180,7 +218,7 @@ def convert_sklearn_gaussian_mixture(scope, operator, container):
         mxlabels = OnnxReduceMax(weighted_log_prob, axes=[1], op_version=opv)
         zeros = OnnxEqual(
             OnnxSub(weighted_log_prob, mxlabels, op_version=opv),
-            np.array([0], dtype=container.dtype),
+            np.array([0], dtype=dtype),
             op_version=opv)
         toint = OnnxCast(zeros, to=onnx_proto.TensorProto.INT64,
                          op_version=opv)


Test patch:
diff --git a/tests/test_op10.py b/tests/test_op10.py
--- a/tests/test_op10.py
+++ b/tests/test_op10.py
@@ -48,9 +48,10 @@ def test_kmeans(self):
     def test_gaussian_mixture(self):
         model, X = fit_classification_model(GaussianMixture(), 3)
         target_opset = 10
-        model_onnx = convert_sklearn(model, "op10",
-                                     [("input", FloatTensorType([None, 3]))],
-                                     target_opset=target_opset)
+        model_onnx = convert_sklearn(
+            model, "op10",
+            [("input", FloatTensorType([None, X.shape[1]]))],
+            target_opset=target_opset)
         self.check_domain(model_onnx, target_opset=target_opset)
 
     @unittest.skipIf(not onnx_built_with_ml(), reason="onnx-ml")
diff --git a/tests/test_sklearn_gaussian_mixture_converter.py b/tests/test_sklearn_gaussian_mixture_converter.py
--- a/tests/test_sklearn_gaussian_mixture_converter.py
+++ b/tests/test_sklearn_gaussian_mixture_converter.py
@@ -67,7 +67,6 @@ def test_model_gaussian_mixture_binary_classification(self):
                     "<= StrictVersion('0.2.1')")
                 self._test_score(model, X, tg)
 
-    @unittest.skipIf(True, reason="Not implemented yet.")
     def test_model_bayesian_mixture_binary_classification(self):
         for cov in ["full", "tied", "diag", "spherical"]:
             with self.subTest(cov=cov):