NumPy memmap in joblib.Parallel — joblib 1.3.0.dev0 documentation

Note

Go to the end to download the full example code

# NumPy memmap in joblib.Parallel

This example illustrates some features enabled by using a memory map (`numpy.memmap`) within `joblib.Parallel`. First, we show that dumping a huge data array ahead of passing it to `joblib.Parallel` speeds up computation. Then, we show the possibility to provide write access to original data.

## Speed up processing of a large data array

> We create a large data array for which the average is computed for several slices.

    import numpy as np

    data = np.random.random((int(1e7),))
    window_size = int(5e5)
    slices = [slice(start, start + window_size)
              for start in range(0, data.size - window_size, int(1e5))]

The `slow_mean` function introduces a `time.sleep()` call to simulate a more expensive computation cost for which parallel computing is beneficial. Parallel may not be beneficial for very fast operation, due to extra overhead (workers creations, communication, etc.).

    import time

    def slow_mean(data, sl):
        """Simulate a time consuming processing."""
        time.sleep(0.01)
        return data[sl].mean()

First, we will evaluate the sequential computing on our problem.

    tic = time.time()
    results = [slow_mean(data, sl) for sl in slices]
    toc = time.time()
    print('\nElapsed time computing the average of couple of slices {:.2f} s'
          .format(toc - tic))

    Elapsed time computing the average of couple of slices 1.03 s

`joblib.Parallel` is used to compute in parallel the average of all slices using 2 workers.

    from joblib import Parallel, delayed

    tic = time.time()
    results = Parallel(n_jobs=2)(delayed(slow_mean)(data, sl) for sl in slices)
    toc = time.time()
    print('\nElapsed time computing the average of couple of slices {:.2f} s'
          .format(toc - tic))

    Elapsed time computing the average of couple of slices 0.67 s

Parallel processing is already faster than the sequential processing. It is also possible to remove a bit of overhead by dumping the `data` array to a memmap and pass the memmap to `joblib.Parallel`.

    import os
    from joblib import dump, load

    folder = './joblib_memmap'
    try:
        os.mkdir(folder)
    except FileExistsError:
        pass

    data_filename_memmap = os.path.join(folder, 'data_memmap')
    dump(data, data_filename_memmap)
    data = load(data_filename_memmap, mmap_mode='r')

    tic = time.time()
    results = Parallel(n_jobs=2)(delayed(slow_mean)(data, sl) for sl in slices)
    toc = time.time()
    print('\nElapsed time computing the average of couple of slices {:.2f} s\n'
          .format(toc - tic))

    Elapsed time computing the average of couple of slices 0.54 s

Therefore, dumping large `data` array ahead of calling `joblib.Parallel` can speed up the processing by removing some overhead.

## Writable memmap for shared memory `joblib.Parallel`

> `slow_mean_write_output` will compute the mean for some given slices as in the previous example. However, the resulting mean will be directly written on the output array.

    def slow_mean_write_output(data, sl, output, idx):
        """Simulate a time consuming processing."""
        time.sleep(0.005)
        res_ = data[sl].mean()
        print("[Worker %d] Mean for slice %d is %f" % (os.getpid(), idx, res_))
        output[idx] = res_

Prepare the folder where the memmap will be dumped.

    output_filename_memmap = os.path.join(folder, 'output_memmap')

Pre-allocate a writable shared memory map as a container for the results of the parallel computation.

    output = np.memmap(output_filename_memmap, dtype=data.dtype,
                       shape=len(slices), mode='w+')

`data` is replaced by its memory mapped version. Note that the buffer has already been dumped in the previous section.

    data = load(data_filename_memmap, mmap_mode='r')

Fork the worker processes to perform computation concurrently

    Parallel(n_jobs=2)(delayed(slow_mean_write_output)(data, sl, output, idx)
                       for idx, sl in enumerate(slices))

    [None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None]

Compare the results from the output buffer with the expected results

    print("\nExpected means computed in the parent process:\n {}"
          .format(np.array(results)))
    print("\nActual means computed by the worker processes:\n {}"
          .format(output))

    Expected means computed in the parent process:
     [0.50017908 0.50039141 0.5001088  0.49982354 0.49989544 0.49991787
     0.49958053 0.49986665 0.50026548 0.50024947 0.50051919 0.50033546
     0.4998536  0.49975586 0.49972224 0.49968148 0.50018648 0.50032584
     0.50000883 0.4999123  0.50003473 0.49980808 0.49951355 0.49997555
     0.50007414 0.4998982  0.49960445 0.50019    0.50037205 0.50034035
     0.50021482 0.50075759 0.50070539 0.50022947 0.50041182 0.50040538
     0.50067041 0.50040977 0.50034548 0.50016992 0.50014481 0.49980217
     0.50001797 0.50034039 0.50022713 0.5003423  0.50000702 0.50009795
     0.49986316 0.49992828 0.50024269 0.5005181  0.50026495 0.50043037
     0.50056342 0.50033946 0.50002878 0.50008725 0.49997608 0.49957387
     0.49940039 0.49977411 0.4995924  0.49944983 0.50019012 0.50026949
     0.49991231 0.50016988 0.50014976 0.4995054  0.49942877 0.49960205
     0.49974909 0.49991553 0.50030638 0.50022496 0.5001178  0.49989686
     0.49993729 0.49991308 0.49996362 0.50005446 0.50021014 0.50025355
     0.50030428 0.50040457 0.50015917 0.49989428 0.49976776 0.49958605
     0.49941402 0.4996179  0.49971412 0.49951242 0.49965507]

    Actual means computed by the worker processes:
     [0.50017908 0.50039141 0.5001088  0.49982354 0.49989544 0.49991787
     0.49958053 0.49986665 0.50026548 0.50024947 0.50051919 0.50033546
     0.4998536  0.49975586 0.49972224 0.49968148 0.50018648 0.50032584
     0.50000883 0.4999123  0.50003473 0.49980808 0.49951355 0.49997555
     0.50007414 0.4998982  0.49960445 0.50019    0.50037205 0.50034035
     0.50021482 0.50075759 0.50070539 0.50022947 0.50041182 0.50040538
     0.50067041 0.50040977 0.50034548 0.50016992 0.50014481 0.49980217
     0.50001797 0.50034039 0.50022713 0.5003423  0.50000702 0.50009795
     0.49986316 0.49992828 0.50024269 0.5005181  0.50026495 0.50043037
     0.50056342 0.50033946 0.50002878 0.50008725 0.49997608 0.49957387
     0.49940039 0.49977411 0.4995924  0.49944983 0.50019012 0.50026949
     0.49991231 0.50016988 0.50014976 0.4995054  0.49942877 0.49960205
     0.49974909 0.49991553 0.50030638 0.50022496 0.5001178  0.49989686
     0.49993729 0.49991308 0.49996362 0.50005446 0.50021014 0.50025355
     0.50030428 0.50040457 0.50015917 0.49989428 0.49976776 0.49958605
     0.49941402 0.4996179  0.49971412 0.49951242 0.49965507]

## Clean-up the memmap

> Remove the different memmap that we created. It might fail in Windows due to file permissions.

    import shutil

    try:
        shutil.rmtree(folder)
    except:  # noqa
        print('Could not clean-up automatically.')

Total running time of the script: ( 0 minutes 2.651 seconds)

Gallery generated by Sphinx-Gallery