Repository: scikit-learn/scikit-learn
Pull request: [MRG+1] Batching in nmf/sparse_dot to prevent MemoryError
#### Reference Issues/PRs
<!--
Example: Fixes #1234. See also #3456.
Please use keywords (e.g., Fixes) to create link to the issues or pull requests
you resolved, so that they will automatically be closed when your pull request
is merged. See https://github.com/blog/1506-closing-issues-via-pull-requests
-->

Fixes https://github.com/scikit-learn/scikit-learn/issues/15242 .


#### What does this implement/fix? Explain your changes.
Batching whilst calculating the sparse dot prevents the allocation of an array with shape (#non-zero elements, rank). I propose to do the batching such that we limit the allocation to an array with shape (#non-zero elements / rank, rank).

#### Any other comments?
Benchmark code to compare the difference in time:
```python
import time

import math
import numpy as np
import scipy.sparse as sp


def batched_simple(W, H, X):
    """Computes np.dot(W, H), only where X is non zero."""
    if sp.issparse(X):
        ii, jj = X.nonzero()
        n_vals = ii.shape[0]
        dot_vals = np.empty(n_vals)
        index = 0
        rank = W.shape[1]
        batch_size = math.floor(n_vals / rank)
        while index < n_vals:
            selector = index + batch_size if index + batch_size <= n_vals else n_vals
            dot_vals[index:selector] = np.multiply(W[ii[index:selector], :], H.T[jj[index:selector], :]).sum(axis=1)
            index = selector

        WH = sp.coo_matrix((dot_vals, (ii, jj)), shape=X.shape)
        return WH.tocsr()
    else:
        return np.dot(W, H)


def original_sparse_dot(W, H, X):
    """Computes np.dot(W, H), only where X is non zero."""
    if sp.issparse(X):
        ii, jj = X.nonzero()
        dot_vals = np.multiply(W[ii, :], H.T[jj, :]).sum(axis=1)
        WH = sp.coo_matrix((dot_vals, (ii, jj)), shape=X.shape)
        return WH.tocsr()
    else:
        return np.dot(W, H)


def generate_values():
    # From the unittests
    n_samples = 1000
    n_features = 50
    n_components = 30
    rng = np.random.mtrand.RandomState(42)
    X = rng.randn(n_samples, n_features)
    np.clip(X, 0, None, out=X)
    X_csr = sp.csr_matrix(X)

    W = np.abs(rng.randn(n_samples, n_components))
    H = np.abs(rng.randn(n_components, n_features))
    return W, H, X_csr


def benchmark():
    W, H, X_csr = generate_values()
    rounds = 100

    def test_func(func, name, rounds, W, H, X_csr):
        start_time = time.clock()
        for i in range(rounds):
            WH = func(W, H, X_csr)
        end_time = time.clock()
        print(name + ": " + str((end_time - start_time) / rounds * 1000) + "ms per loop.")

    test_func(original_sparse_dot, "original version", rounds, W, H, X_csr)
    test_func(batched_simple, "batched_simple", rounds, W, H, X_csr)


if __name__ == '__main__':
    benchmark()
```
Output:
```
original version: 11.216264ms per loop.
batched_simple: 4.231038999999999ms per loop.
```
I think this to be a strange result: I would have expected one big batch to be better. Discussion welcome :-)


Problem statement:
MemoryError when using NMF's multiplicative update solver
#### Description
MemoryError when using NMF's multiplicative update solver with method="Kullback_leibler". The program tries to assign a matrix with size (# nonzero elements, rank) in function _special_sparse_dot (in decomposition/nmf.py) and sums the columns. On my machine, this results in a MemoryError. The dot product, however, could be calculated element-wise, requiring less memory and allowing for a larger matrix to be processed
#### Steps/Code to Reproduce
```python
from scipy.sparse import random
from sklearn.decomposition import NMF

matrix = random(1500, 5000, density=1)
nmf = NMF(n_components=3000, random_state=1,
          alpha=.1, l1_ratio=.5, verbose=True, beta_loss="kullback-leibler", solver="mu")
fit = nmf.fit_transform(matrix)
```

#### Expected Results
No error is thrown; dot product is calculated without extreme memory requirements. A potential fix could be to calculate the dot product element wise:
```python
def _special_sparse_dot(W, H, X):
    """Computes np.dot(W, H), only where X is non zero."""
    if sp.issparse(X):
        ii, jj = X.nonzero()
        dot_vals = []
        for i, _ in enumerate(ii):
            wv = W[ii[i], :]
            hv = H.T[jj[i], :]
            dot_vals.append(np.dot(wv, hv))

        # dot_vals = np.multiply(W[ii, :], H.T[jj, :]).sum(axis=1)
        WH = sp.coo_matrix((dot_vals, (ii, jj)), shape=X.shape)
        return WH.tocsr()
    else:
        return np.dot(W, H)

```

#### Actual Results
<!-- Please paste or specifically describe the actual output or traceback. -->
```python 
Traceback (most recent call last):
  File "<input>", line 1, in <module>
  File "C:\Program Files\JetBrains\PyCharm 2019.1.3\helpers\pydev\_pydev_bundle\pydev_umd.py", line 197, in runfile
    pydev_imports.execfile(filename, global_vars, local_vars)  # execute the script
  File "C:\Program Files\JetBrains\PyCharm 2019.1.3\helpers\pydev\_pydev_imps\_pydev_execfile.py", line 18, in execfile
    exec(compile(contents+"\n", file, 'exec'), glob, loc)
  File "C:/Users/<user>/.PyCharm2019.1/config/scratches/scratch.py", line 9, in <module>
    fit = nmf.fit_transform(matrix)
  File "C:\Workdir\programs\miniconda\envs\monty\lib\site-packages\sklearn\decomposition\nmf.py", line 1289, in fit_transform
    shuffle=self.shuffle)
  File "C:\Workdir\programs\miniconda\envs\monty\lib\site-packages\sklearn\decomposition\nmf.py", line 1077, in non_negative_factorization
    verbose)
  File "C:\Workdir\programs\miniconda\envs\monty\lib\site-packages\sklearn\decomposition\nmf.py", line 800, in _fit_multiplicative_update
    error_at_init = _beta_divergence(X, W, H, beta_loss, square_root=True)
  File "C:\Workdir\programs\miniconda\envs\monty\lib\site-packages\sklearn\decomposition\nmf.py", line 119, in _beta_divergence
    WH_data = _special_sparse_dot(W, H, X).data
  File "C:\Workdir\programs\miniconda\envs\monty\lib\site-packages\sklearn\decomposition\nmf.py", line 187, in _special_sparse_dot
    dot_vals = np.multiply(W[ii, :], H.T[jj, :]).sum(axis=1)
numpy.core._exceptions.MemoryError: Unable to allocate array with shape (75000000, 300) and data type float64
```
#### Versions
System:
    python: 3.6.8 |Anaconda, Inc.| (default, Feb 21 2019, 18:30:04) [MSC v.1916 64 bit (AMD64)]
executable: C:\Workdir\programs\miniconda\envs\monty\python.exe
   machine: Windows-10-10.0.17763-SP0
Python deps:
       pip: 19.1.1
setuptools: 41.0.1
   sklearn: 0.21.3
     numpy: 1.17.0
     scipy: 1.3.0
    Cython: None
    pandas: 0.25.0

<!-- Thanks for contributing! -->

Hints from issue discussion:
That is indeed a bug, thanks for the report and the short reproducing example.
PR most welcome !

The loop will be very slow, so we might need to do some batching or some cython here.

Commit message: Avoid multiple statements on one line

Patch to review:
diff --git a/sklearn/decomposition/nmf.py b/sklearn/decomposition/nmf.py
index 0cf663e123861..d0a00bf5e8018 100644
--- a/sklearn/decomposition/nmf.py
+++ b/sklearn/decomposition/nmf.py
@@ -6,20 +6,19 @@
 #         Tom Dupre la Tour
 # License: BSD 3 clause
 
-from math import sqrt
-import warnings
 import numbers
-import time
-
 import numpy as np
 import scipy.sparse as sp
+import time
+import warnings
+from math import sqrt, floor
 
+from .cdnmf_fast import _update_cdnmf_fast
 from ..base import BaseEstimator, TransformerMixin
+from ..exceptions import ConvergenceWarning
 from ..utils import check_random_state, check_array
 from ..utils.extmath import randomized_svd, safe_sparse_dot, squared_norm
 from ..utils.validation import check_is_fitted, check_non_negative
-from ..exceptions import ConvergenceWarning
-from .cdnmf_fast import _update_cdnmf_fast
 
 EPSILON = np.finfo(np.float32).eps
 
@@ -170,7 +169,19 @@ def _special_sparse_dot(W, H, X):
     """Computes np.dot(W, H), only where X is non zero."""
     if sp.issparse(X):
         ii, jj = X.nonzero()
-        dot_vals = np.multiply(W[ii, :], H.T[jj, :]).sum(axis=1)
+        n_vals = ii.shape[0]
+        dot_vals = np.empty(n_vals)
+        i = 0
+        rank = W.shape[1]
+        batch_size = floor(n_vals / rank)
+        while i < n_vals:
+            s = i + batch_size
+            if s > n_vals:
+                s = n_vals
+            dot_vals[i:s] = np.multiply(W[ii[i:s], :],
+                                        H.T[jj[i:s], :]).sum(axis=1)
+            i = s
+
         WH = sp.coo_matrix((dot_vals, (ii, jj)), shape=X.shape)
         return WH.tocsr()
     else:


Code review comments:
diff_hunk: @@ -170,7 +169,19 @@ def _special_sparse_dot(W, H, X):
     """Computes np.dot(W, H), only where X is non zero."""
     if sp.issparse(X):
         ii, jj = X.nonzero()
-        dot_vals = np.multiply(W[ii, :], H.T[jj, :]).sum(axis=1)
+        n_vals = ii.shape[0]
+        dot_vals = np.empty(n_vals)
+        i = 0
+        rank = W.shape[1]
+        batch_size = floor(n_vals / rank)
+        while i < n_vals:
+            s = i + batch_size
+            if s > n_vals:
path: sklearn/decomposition/nmf.py
text: @user1:
Note that slices are automatically cropped by python, so we do not need a special case for the last one:
```py
print(list(range(10))[7:13])
# [7, 8, 9]
```

diff_hunk: @@ -170,7 +169,19 @@ def _special_sparse_dot(W, H, X):
     """Computes np.dot(W, H), only where X is non zero."""
     if sp.issparse(X):
         ii, jj = X.nonzero()
-        dot_vals = np.multiply(W[ii, :], H.T[jj, :]).sum(axis=1)
+        n_vals = ii.shape[0]
+        dot_vals = np.empty(n_vals)
+        i = 0
+        rank = W.shape[1]
path: sklearn/decomposition/nmf.py
text: @user1:
I would rather call it `n_components` for consistency.

diff_hunk: @@ -170,7 +169,19 @@ def _special_sparse_dot(W, H, X):
     """Computes np.dot(W, H), only where X is non zero."""
     if sp.issparse(X):
         ii, jj = X.nonzero()
-        dot_vals = np.multiply(W[ii, :], H.T[jj, :]).sum(axis=1)
+        n_vals = ii.shape[0]
+        dot_vals = np.empty(n_vals)
+        i = 0
+        rank = W.shape[1]
+        batch_size = floor(n_vals / rank)
path: sklearn/decomposition/nmf.py
text: @user1:
This can be zero if `n_vals < rank`, leading to an infinite loop.
You can also use `//` instead of `floor`.
What do you think of `batch_size = max(rank, n_vals // rank)` ?

diff_hunk: @@ -170,7 +169,19 @@ def _special_sparse_dot(W, H, X):
     """Computes np.dot(W, H), only where X is non zero."""
     if sp.issparse(X):
         ii, jj = X.nonzero()
-        dot_vals = np.multiply(W[ii, :], H.T[jj, :]).sum(axis=1)
+        n_vals = ii.shape[0]
+        dot_vals = np.empty(n_vals)
+        i = 0
+        rank = W.shape[1]
+        batch_size = floor(n_vals / rank)
+        while i < n_vals:
path: sklearn/decomposition/nmf.py
text: @user1:
I tend to prefer `for` loops, to avoid infinite loops:
```py
for start in range(0, n_vals, batch_size):
    batch = slice(start, start + batch_size)
    dot_vals[batch] = np.multiply(W[ii[batch], :],
                                  H.T[jj[batch], :]).sum(axis=1)
```

@author:
Thanks for taking the time to review my code! I learned a lot from this!