-
-
Notifications
You must be signed in to change notification settings - Fork 26k
ENH avoid np.square(X) in enet_coordinate_descent to save memory #31665
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
ENH avoid np.square(X) in enet_coordinate_descent to save memory #31665
Conversation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Do you have a quick memory benchmark of einsum
vs np.square(X).sum(axis=0)
?
import numpy as np
rng = np.random.default_rng(42)
X = rng.standard_normal((100, 10_000))
print(f"X allocates {X.nbytes * 1e-6} MB of memory")
# X allocates 8.0 MB of memory
%timeit np.square(X).sum(axis=0)
# 1.09 ms ± 53.2 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
%timeit np.einsum("ij,ij->j", X, X)
# 631 μs ± 29.3 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
%load_ext memory_profiler
%memit np.square(X).sum(axis=0)
peak memory: 87.97 MiB, increment: 7.98 MiB
In [6]: %memit np.einsum("ij,ij->j", X, X)
peak memory: 88.34 MiB, increment: 0.03 MiB Similar with transposed X. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM. Thanks @lorentzenchr
Reference Issues/PRs
None
What does this implement/fix? Explain your changes.
This PR replaces
np.square(X).sum(axis=0)
bynp.einsum("ij,ij->j", X, X)
to avoid memory allocation of the size ofX
(usually the largest object).Any other comments?
This also improves timing a bit.
We might even consider to write the loop explicitly like in (
scikit-learn/sklearn/metrics/_pairwise_distances_reduction/_base.pyx.tp
Lines 20 to 26 in 20b8d0b