Data multiprocessing in Python

Quick guide to making your GPU inference (or any data-processing pipeline) faster through multiprocessing.

A basic data-processing pipeline has three components:

1. Input reader: loads and prepares the data. In PyTorch, the DataLoader already handles this in parallel, but other inputs may need their own reader processes.
2. Compute: runs inference or tokenization on one device (typically a GPU), or on multiple devices in parallel. Each device has its own process and model copy.
3. Output processing: receives completed results through a queue and works on them while compute continues. In our examples this handles writing files to disk, but it could be any CPU-side processing.

Each component can run in parallel in a producer-consumer pattern.

input reader(s)  ->  compute process(es)  --(bounded) queue-->  output worker(s)
   DataLoader          one per device                             save/process

Files

• inference.py: commented version with progress logging.
• inference_tqdm.py: the same version using tqdm, a library that displays a live progress bar.
• inference_tqdm_multiprocess.py: tqdm-only version with one fixed progress line per process (experts-only).
• inference_compact.py: compact version of the same pipeline.
• inference_multigpu.py: use also multiple GPUs in parallel.

Setup

uv venv .env
source .env/bin/activate
uv pip install -r requirements.txt

Run

source .env/bin/activate
python inference.py --samples 1_000_000
python inference_tqdm.py --samples 1_000_000
python inference_tqdm_multiprocess.py --samples 1_000_000
python inference_compact.py --samples 1_000_000
python inference_multigpu.py --samples 1_000_000

The useful tuning knobs are:

• --batch-size
• --loader-workers
• --writers
• --queue-size.

More processes are not always faster; tune them for the dataset, CPU count, and storage system.

inference_multigpu.py uses all visible GPUs by default. Pass --gpus 0,1 to select devices or --gpus cpu for a CPU fallback.