`sklearn.dataset.make_blobs` takes too much RAM

[NightMachinery]

Describe the bug

sklearn.dataset.make_blobs consumes an unusual amount of memory.

Steps/Code to Reproduce

command time -f '%M' python -c '
from sklearn import datasets

blobs_opts = {
    "n_samples": 10**4,
    "n_features": 10**4,
    "centers": 10,
    "random_state": 10
}
X, y = datasets.make_blobs(**blobs_opts)
None
'

Expected Results

The max memory consumption should be around 10**(4+4)*8=800000000 bytes = 763MB.

Actual Results

The max memory consumption is 2426524 kilobytes = 2.4GB. (The time command returns the result in kilobytes.)

Versions

System:
    python: 3.8.10 | packaged by conda-forge | (default, May 11 2021, 07:01:05)  [GCC 9.3.0]
executable: /root/miniconda3/bin/python
   machine: Linux-5.4.144+-x86_64-with-glibc2.10

Python dependencies:
          pip: 21.3.1
   setuptools: 59.8.0
      sklearn: 1.0.2
        numpy: 1.21.5
        scipy: 1.7.3
       Cython: None
       pandas: 1.3.5
   matplotlib: 3.4.3
       joblib: 1.1.0
threadpoolctl: 3.0.0

Built with OpenMP: True

[NightMachinery]

Strangely enough, when I use:

command time -f 'Max memory: %M KB' python -c '
import sklearn
sklearn.show_versions()
from sklearn import datasets

blobs_opts = {
    "n_samples": 10**4,
    "n_features": 10**4,
    "centers": 10,
    "random_state": 10
}
X, y = datasets.make_blobs(**blobs_opts)
None
'

I get 1709112 KB = 1.7GB. The results are robust across multiple runs.

[MaxwellLZH]

Running with memory_profiler on MacOS and latest scikit version gives the following result:

Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
    19     79.3 MiB     79.3 MiB           1   @profile
    20                                         def run():
    21     79.3 MiB      0.0 MiB           1       blobs_opts = {
    22     79.3 MiB      0.0 MiB           1           "n_samples": 10**4,
    23     79.3 MiB      0.0 MiB           1           "n_features": 10**4,
    24     79.3 MiB      0.0 MiB           1           "centers": 10,
    25     79.3 MiB      0.0 MiB           1           "random_state": 10
    26                                             }
    27   1605.4 MiB   1526.1 MiB           1       X, y = datasets.make_blobs(**blobs_opts)

[glemaitre]

Basically, we create a temporary array and then concatenate:

scikit-learn/sklearn/datasets/_samples_generator.py

Lines 933 to 949 in ddbf5b3

	X = []
	y = []
	if isinstance(n_samples, Iterable):
	n_samples_per_center = n_samples
	else:
	n_samples_per_center = [int(n_samples // n_centers)] * n_centers
	for i in range(n_samples % n_centers):
	n_samples_per_center[i] += 1
	for i, (n, std) in enumerate(zip(n_samples_per_center, cluster_std)):
	X.append(generator.normal(loc=centers[i], scale=std, size=(n, n_features)))
	y += [i] * n
	X = np.concatenate(X)
	y = np.array(y)

We could avoid the temporary arrays by generating the random data directly within a pre-allocated empty NumPy array.

The max memory consumption should be around 10**(4+4)*8=800000000 bytes = 763MB.

In your computation, you are missing that we need to allocate y and the centers as well.

[glemaitre]

A potential alternative implementation would be:

    cum_sum_n_samples = np.cumsum(n_samples_per_center)
    X = np.empty(shape=(sum(n_samples_per_center), n_features), dtype=np.float64)
    y = np.empty(shape=(sum(n_samples_per_center),), dtype=np.int64)

    for i, (n, std) in enumerate(zip(n_samples_per_center, cluster_std)):
        start_idx = cum_sum_n_samples[i - 1] if i > 0 else 0
        end_idx = cum_sum_n_samples[i]
        X[start_idx:end_idx] = generator.normal(
            loc=centers[i], scale=std, size=(n, n_features)
        )
        y[start_idx:end_idx] = i

    if shuffle:
        X, y = util_shuffle((X, y), random_state=generator)

I am getting the following memory increase:

In [3]: %%memit
   ...: from sklearn import datasets
   ...: 
   ...: blobs_opts = {
   ...:     "n_samples": 10**4,
   ...:     "n_features": 10**4,
   ...:     "centers": 10,
   ...:     "random_state": 10
   ...: }
   ...: X, y = datasets.make_blobs(**blobs_opts)
   ...: 
   ...: 

peak memory: 1705.98 MiB, increment: 762.95 MiB

which is closer to what you would expect.

[NightMachinery]

Yes, the increment is totally reasonable, but why does the peak memory have an extra GB? I generally don’t like memory_profiler (which the memit magic uses under the hood); its results are weird and dependent on the interpreter state (e.g., if it has memory claimed from the OS that it has not used). Using the time command is completely reliable, on the other hand.

If the total memory usage should be around 1GB, it’s very reasonable. Can you merge your improvement in?

BTW, can the dtype of the generated data be specified in the API? I am trying to compare some algorithms’ memory usages, and I use float32, which means that the memory overhead of first generating the data in float64 covers up the memory inefficiencies in the algorithms I want to benchmark.
The RAPIDS version of this function does accept the dtype, so I think it’s a reasonable API choice. I can open another issue if it’s out of scope for this one.

PS: I did some quick calculations before, I think y and the centers consume negligible (less than 10MB) memory.

[glemaitre]

BTW, can the dtype of the generated data be specified in the API? I am trying to compare some algorithms’ memory usages, and I use float32, which means that the memory overhead of first generating the data in float64 covers up the memory inefficiencies in the algorithms I want to benchmark.
The RAPIDS version of this function does accept the dtype, so I think it’s a reasonable API choice. I can open another issue if it’s out of scope for this one.

Since it might be more and more common to use np.float32 as dataset, we could provide such option.

[glemaitre]

Can you merge your improvement in?

Do you want to propose a PR to fix it? I could then provide a first approval on the 2 approvals requested.

[NightMachinery]

Can you merge your improvement in?

Do you want to propose a PR to fix it? I could then provide a first approval on the 2 approvals requested.

Okay, I’ll open a PR. Should I put both the memory fix and the dtype option into a single PR?

---

I also have another unrelated idea: It would be great if the dataset generation APIs had an API to write big datasets (that don’t fit in the memory) to disk. In the case of make_blobs, I think this can be achieved by first generating the centers, and then generating a part of the big dataset, concatenating it to the destination file, freeing up its memory, and repeating till all the parts are done.

(This feature is useful when testing out-of-core algorithms.)

Do you think that makes sense? Should I open a PR for that as well?

[NightMachinery]

Another matter came to my mind; is there a way to add a unit test to ensure that there wouldn’t be a regression in memory usage in the future? Would it be a good idea to add such a test in my PR?

[glemaitre]

Okay, I’ll open a PR. Should I put both the memory fix and the dtype option into a single PR?

It would be best to split. We might want to discuss a bit more about the dtype because it should also be included in other generators I think.

I also have another unrelated idea: It would be great if the dataset generation APIs had an API to write big datasets (that don’t fit in the memory) to disk. In the case of make_blobs, I think this can be achieved by first generating the centers, and then generating a part of the big dataset, concatenating it to the destination file, freeing up its memory, and repeating till all the parts are done.

I don't know if it should be the job of scikit-learn here. I would think that dask and dask-ml would be better suited in these cases.

Another matter came to my mind; is there a way to add a unit test to ensure that there wouldn’t be a regression in memory usage in the future? Would it be a good idea to add such a test in my PR?

We can discuss it in the PR when it will be reviewed. It would be great to indeed have a non-regression test as you mentioned. When it comes to performance as memory consumption this is a bit more complicated to check.