Skip to content

Sharing memory among processes

The Python multiprocessing.shared_memory module enables to share memory among processes.

The shared_memory_example.py script demonstrate this.

Note

You need to install numpy to run this example, since it is used to save the resulting masks stored in the shared memory.

It will use shared_memory_module.py to create the segmentation and the shared memory holding the resulting masks.

This module defines two functions:

  • a segment function which uses the segment function of example_module.py and creates a NumPy array backed by a shared memory to store the resulting masks,

  • a clean function to clean up, free and release the shared memory block.

import numpy as np
from pathlib import Path
from multiprocessing import shared_memory

import example_module

shm: shared_memory.SharedMemory | None = None


def segment(imagePath: Path | str):
    global shm
    # Segment the image with example_module.py
    masks, flows, styles, diams = example_module.segment(imagePath)
    # Create the shared memory
    shm = shared_memory.SharedMemory(create=True, size=masks.nbytes)
    # Create a NumPy array backed by shared memory
    masksShared = np.ndarray(masks.shape, dtype=masks.dtype, buffer=shm.buf)
    # Copy the masks into the shared memory
    masksShared[:] = masks[:]
    # Return the shape, dtype and shared memory name to recreate the numpy array on the other side
    return masks.shape, masks.dtype, shm.name


def clean():
    global shm
    if shm is None:
        return
    # Clean up the shared memory in this process
    shm.close()
    # Free and release the shared memory block
    shm.unlink()

The shared_memory_example.py script creates an environment using the initialization function from getting_started.py.

shared_memory_example.py 

from multiprocessing import shared_memory
import numpy as np

import getting_started

# Create a Conda environment from getting_started.py
imagePath, segmentationPath, env = getting_started.initialize()

Then, it imports shared_memory_module.py to perform a cellpose segmentation, and creates a shared memory for the resulting masks.

# Import shared_memory_module in the environment
sharedMemoryModule = env.importModule("shared_memory_module.py")
# run env.execute(module_name, function_name, args)
diameters, masksShape, masksDtype, shmName = sharedMemoryModule.segment(imagePath)

This shared memory can now be used in the main process, for example to save the masks as a numpy binary file:

# Save the segmentation from the shared memory
shm = shared_memory.SharedMemory(name=shmName)
masks = np.ndarray(masksShape, dtype=masksDtype, buffer=shm.buf)
segmentationPath = imagePath.parent / f"{imagePath.stem}_segmentation.bin"
masks.tofile(segmentationPath)

print(f"Found diameters of {diameters} pixels.")

# Clean up and exit the environment
env.exit()

Note

Changes to the shared memory made by one process will be refelcted in the other process. You can update the memory from both processes and perform more sofisticated operations.