使用torch.multiprocessing.Queue未获性能提升，求解原因

I have a training system composed of multiple processes that exchange tensors, lists of tensors, or dictionaries. I expected using torch.multiprocessing's shared memory to speed up the workflow, but enabling it actually made things slower. Here's my test script:

import sys
import time
import torch
from torch.multiprocessing import Process as TorchProcess
from torch.multiprocessing import Queue as TorchQueue

q = TorchQueue()

def torch_shared_mem_process():
    counter = 0
    while True:
        data = q.get()
        counter += 1
        if data is None:
            return
        print('Received data:', len(data), data, counter)

def test_mem_share(share_memory):
    p = TorchProcess(target=torch_shared_mem_process)
    p.start()
    def sample_data():
        return torch.rand([1000, 128, 72, 3], dtype=torch.float)
    start = time.time()
    n = 50
    for i in range(n):
        data = sample_data()
        for data_item in data:
            if share_memory:
                data_item.share_memory_()
        q.put(data)
        print(f'Progress {i}/{n}')
    q.put(None)
    p.join()
    print(f'Finished sending {n} tensor lists!')
    took_seconds = time.time() - start
    return took_seconds

def main():
    no_shared_memory = test_mem_share(share_memory=False)
    with_shared_memory = test_mem_share(share_memory=True)
    print(f'Took {no_shared_memory:.1f} s without shared memory.')
    print(f'Took {with_shared_memory:.1f} s with shared memory.')

if __name__ == '__main__':
    sys.exit(main())

My test results show the shared memory version is slower:

Took 10.2 s without shared memory.
Took 11.7 s with shared memory.

Did I misunderstand how torch.multiprocessing.Queue works?

Great question—let's break down why you're seeing this unexpected slowdown, and how to fix it.

1. You're calling share_memory_() inefficiently #

The biggest issue in your code is that you're looping through every slice of your large tensor and calling share_memory_() on each one individually. Each call to share_memory_() involves:

• Allocating shared memory from the OS
• Copying the tensor's data into that shared region
• Updating the tensor's storage to point to the shared memory

Doing this 1000 times per iteration (for each data_item in your [1000, 128, 72, 3] tensor) adds massive unnecessary overhead. Instead, you should call share_memory_() once on the entire tensor:

# Replace this loop:
for data_item in data:
    if share_memory:
        data_item.share_memory_()

# With this single call:
if share_memory:
    data.share_memory_()

This reduces the number of OS-level shared memory operations from 1000 per tensor to 1, which will drastically cut the overhead.

2. torch.multiprocessing.Queue already optimizes tensor transfers #

You might not realize this, but torch.multiprocessing.Queue uses a custom pickler that's optimized for PyTorch tensors. Even without manually enabling shared memory, it avoids full data copies for tensors in many cases (especially when sending large tensors between processes).

When you don't call share_memory_(), the queue uses a fast, zero-copy serialization method for tensors. The overhead here is minimal compared to thousands of share_memory_() calls.

3. Your test doesn't leverage shared memory's benefits #

Shared memory shines when processes need to repeatedly access the same tensor data without copying it. In your test, the receiver process only prints metadata about the tensor (length, object reference) and doesn't actually access the tensor's underlying data.

Without accessing the data, you don't get the benefit of avoiding copies to the receiver's memory space—you just pay the cost of setting up shared memory without reaping any rewards. If your real workflow involves the receiver process using the tensor data (e.g., running computations on it), the shared memory version would pull ahead.

Fixing your test script #

Here's the modified version of your test_mem_share function that uses shared memory correctly:

def test_mem_share(share_memory):
    p = TorchProcess(target=torch_shared_mem_process)
    p.start()
    def sample_data():
        return torch.rand([1000, 128, 72, 3], dtype=torch.float)
    start = time.time()
    n = 50
    for i in range(n):
        data = sample_data()
        if share_memory:
            # Apply shared memory to the entire tensor, not slices
            data.share_memory_()
        q.put(data)
        print(f'Progress {i}/{n}')
    q.put(None)
    p.join()
    print(f'Finished sending {n} tensor lists!')
    took_seconds = time.time() - start
    return took_seconds

With this change, you should see the shared memory version perform on par or better than the non-shared version, especially if you modify the receiver process to actually use the tensor data (e.g., compute data.sum()).

内容的提问来源于stack exchange，提问作者Aleksei Petrenko