使用torch.multiprocessing.Queue获取Tensor触发RuntimeError，如何实现多进程获取？

你遇到的RuntimeError是因为PyTorch Tensor默认依赖共享内存机制实现进程间数据交互，但直接使用标准库的multiprocessing.Queue传递时，共享内存的事件管理出现异常。而numpy数组采用值传递，不会触发这类问题。以下是几种可行的解决方法：

方法1：使用PyTorch适配的Queue #

PyTorch对torch.multiprocessing.Queue做了专门优化，原生支持Tensor的进程间传递，替换标准库队列即可解决问题：

import torch
import torch.multiprocessing as mp
import numpy as np

def solve(queue):
    for i in range(3):
        queue.put(torch.ones((2, 2)))

def func():
    jobs = []
    # 改用PyTorch提供的Queue
    queue = mp.Queue()

    for i in range(5):
        p = mp.Process(target=solve, args=(queue,))
        jobs.append(p)
        p.start()

    for j in jobs:
        j.join()

    print('end')
    out = queue.get()
    print(out)

if __name__ == '__main__':
    func()

方法2：手动将Tensor转为共享内存模式 #

如果必须使用标准库的multiprocessing.Queue，可以手动调用share_memory_()方法将Tensor转为共享内存模式：

import torch
import multiprocessing as mp
import numpy as np

def solve(queue):
    for i in range(3):
        tensor = torch.ones((2, 2))
        # 开启共享内存
        tensor.share_memory_()
        queue.put(tensor)

def func():
    jobs = []
    queue = mp.Queue()

    for i in range(5):
        p = mp.Process(target=solve, args=(queue,))
        jobs.append(p)
        p.start()

    for j in jobs:
        j.join()

    print('end')
    out = queue.get()
    print(out)

if __name__ == '__main__':
    func()

方法3：转成numpy数组传递后再转回Tensor #

这种方式和你测试numpy数组可行的逻辑一致，通过值传递避免共享内存问题：

import torch
import multiprocessing as mp
import numpy as np

def solve(queue):
    for i in range(3):
        tensor = torch.ones((2, 2))
        # 转为numpy数组传递
        queue.put(tensor.numpy())

def func():
    jobs = []
    queue = mp.Queue()

    for i in range(5):
        p = mp.Process(target=solve, args=(queue,))
        jobs.append(p)
        p.start()

    for j in jobs:
        j.join()

    print('end')
    out_np = queue.get()
    # 从numpy数组转回Tensor
    out = torch.from_numpy(out_np)
    print(out)

if __name__ == '__main__':
    func()

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