PERF float32 propagation in GaussianMixture

[ogrisel]

This is a draft PR to explore how much work is need to ensure floating-point precision propagation in GaussianMixture as part of the discussion in #30382.

There are still some failing tests to fix.

[github-actions]

✔️ Linting Passed

All linting checks passed. Your pull request is in excellent shape! ☀️

_{Generated for commit: 8fd0b71. Link to the linter CI: here}

[OmarManzoor]

@ogrisel I added a small adjustment and the float32 tests seem to pass now. Do we have any way to invoke them in the CI as well?
Edit: I also updated two tests and adjusted the tolerance which checks for normalized weights because in the test test_gaussian_mixture_fit_predict it was failing due to the weights summing to around 1.0000001

[ogrisel]

Do we have any way to invoke them in the CI as well?

Yes, by pushing a commit with the [float32] marker in the commit message:

https://scikit-learn.org/stable/developers/contributing.html#commit-message-markers

[OmarManzoor]

Yes, by pushing a commit with the [float32] marker in the commit message:

https://scikit-learn.org/stable/developers/contributing.html#commit-message-markers

Okay nice let me see if the CI passes first before trying out float32 because for the initial pipelines the np.log did not accept a dtype argument.

[ogrisel]

I confirm that the tests now pass locally. Let me add a changelog entry.

[ogrisel]

Let me do a quick benchmark to see the impact of this change.

[ogrisel]

On my Apple M1 with accelerate (mamba install "libblas=*=*accelerate") I get the following numbers:

• on this branch: 7.8 s with a 1 GB memory usage peak
• on main: 23.1 s with a 2 GB memory usage peak

So a nice ~3X speed up and a 2X memory efficiency improvement.

Benchmark script:

from sklearn.mixture import GaussianMixture
from sklearn.datasets import make_blobs
import numpy as np
from time import perf_counter

X, y = make_blobs(n_samples=int(1e6), n_features=50, centers=10, random_state=0)
X = X.astype(np.float32)

tic = perf_counter()
GaussianMixture(n_components=10, covariance_type="full", random_state=0).fit(X)
toc = perf_counter()
print(f"Elapsed time: {toc - tic:.1f} s")

I used the scalene profiler to report memory usage (scalene bench_gmm.py with the script above to reproduce).

[ogrisel]

I also tried with OpenBLAS (mamba install "libblas=*=*openblas") and got a 1.9x speed up instead of nearly 3x. I suspect that Accelerate can use the Apple M1 GPU for some operations when working with float32 data while OpenBLAS is limited to the CPU vector instruction set.

Memory usage is the same for OpenBLAS and Accelerate.

[OmarManzoor]

I also tried with OpenBLAS (mamba install "libblas=*=*openblas") and got a 1.9x speed up instead of nearly 3x. I suspect that Accelerate can use the Apple M1 GPU for some operations when working with float32 data while OpenBLAS is limited to the CPU vector instruction set.

Even then there is a nice speedup.

[ogrisel]

Once this is merged, I think this estimator would be a nice candidate for array API support.

[mekleo]

On my Apple M1 with accelerate (mamba install "libblas=*=*accelerate") I get the following numbers:

• on this branch: 7.8 s with a 1 GB memory usage peak
• on main: 23.1 s with a 2 GB memory usage peak

So a nice ~3X speed up and a 2X memory efficiency improvement.

Benchmark script:

from sklearn.mixture import GaussianMixture
from sklearn.datasets import make_blobs
import numpy as np
from time import perf_counter

X, y = make_blobs(n_samples=int(1e6), n_features=50, centers=10, random_state=0)
X = X.astype(np.float32)

tic = perf_counter()
GaussianMixture(n_components=10, covariance_type="full", random_state=0).fit(X)
toc = perf_counter()
print(f"Elapsed time: {toc - tic:.1f} s")

I used the scalene profiler to report memory usage (scalene bench_gmm.py with the script above to reproduce).

Great work! Beyond performance enhancement, in your tests, do you obtain identical results after downcasting X to np.float32?

[ogrisel]

Great work! Beyond performance enhancement, in your tests, do you obtain identical results after downcasting X to np.float32?

The existing tests still pass when we run them with float32 (up to some tolerance level that can depend on the choice of the dtype).

[betatim]

Looks reasonable to me. A few little comments.

One question: is there a unit test that checks that the dtype of gmm.predict(X) matches the dtype of X or should it match the dtype of X that was used for fitting? Or do we not need that kind of test for some reason?

[ogrisel]

One question: is there a unit test that checks that the dtype of gmm.predict(X) matches the dtype of X or should it match the dtype of X that was used for fitting?

gmm.predict(X) returns the integer index of the component with the highest likelihood. So it is unrelated to the dtype of X.

[lesteve]

I am not sure it matters that much, but it seems like resp is always going to be float64 (apparently .uniform does not let you specify a dtype). In particular the .astype(X.dtype) is not going to affect resp.dtype?

[betatim]

Indeed! TIL. If resp.dtype is not float32 then dividing by a float32 won't change the dtype of it :-/

I guess we'd have to add a resp = resp.astype(X.dtype) and maybe can remove the astype from resp /= resp.sum(axis=1)[:, np.newaxis].astype(X.dtype)?

[OmarManzoor]

Yes this was a valid point. Thanks for catching it @lesteve
I updated this and also added global_dtype in a test that uses this particular initialization.

[ogrisel]

Thanks, @lesteve and @OmarManzoor for the latest round of fixes / testing.

[lesteve]

I am not sure this makes a big difference but we may as well avoid a copy if X.dtype is float64 (which I guess is the most common use case)?

Suggested change

	resp = random_state.uniform(size=(n_samples, self.n_components)).astype(
	X.dtype
	)
	resp = np.asarray(
	random_state.uniform(size=(n_samples, self.n_components)), dtype=X.dtype
	)

[lesteve]

Thanks, I have enabled auto-merge!

[lucyleeow]

@ogrisel dtype here seems not to be used? Was it meant to be?

Ref: #31432