用multiprocessing.Queue自己编写进程池实现DataLoader(父子进程通信)

原创

已于 2022-01-22 16:42:56 修改 · 1k 阅读

2 ·

CC 4.0 BY-SA版权

文章标签：

#python #多进程 #深度学习 #算法

于 2022-01-19 12:57:03 首次发布

本文介绍如何利用multiprocessing库中的Queue实现父子进程间的通信，并基于此原理自定义DataLoader，用于数据预处理。此外，还探讨了解决多进程环境下主进程被终止后遗留子进程的问题。

1. 用Queue写进程池从而实现父子进程通信

2. 自己编写DataLoader

3. 关于多进程杀掉主进程会有子进程残留的问题

multiprocess.Queue是多进程安全的队列，可以使用Queue实现多进程之间的数据传递, 底层使用管道pipe，同步信号量和互斥锁实现。

multiprocess库提供了pool进程池，直接实现了多进程之间的通信，pool的使用场景也很多，这里不做介绍。但pool不能实现父子进程之间的通信，要想实现父子通信，需要自己用Queue写一个进程池，通过创建子队列和父队列来进行父子通信。这个用途就更广泛了，比如想要开多个子进程帮你处理一些数据或文件，最后把结果都收回到主进程，这也就是自己编写DataLoader，很多实际项目中需要自己写一个数据预处理代码，所以要自己重写DataLoader。

本文先讲如何实现父子进程通信，之后就很容易实现DataLoader了。

1. 用Queue写进程池从而实现父子进程通信

思路其实很简单，先分别创建子队列父队列，然后在主进程中把文件放入子队列，多个子进程各自从子队列中取文件并进行处理，再把各自处理好的数据放入父队列，主进程等所有文件处理完后从父队列取处理好的数据。

这里比较需要注意的是用None作为退出信号，因为如果是用Queue的get()函数（默认block=False）则会在队列为空时阻塞等待而不退出；如果用get_nowait()也即get(block=True)，则会在队列为空时报出异常；设置成当队列为空则退出也是不合理的。因此可以考虑用None作为退出信号，由主进程在文件都放入子队列后发出。

这里省略了具体处理data的函数。由于比较容易理解所以直接上代码。

from multiprocessing import Process, Queue


def data_preprocess(file, args):
    # preprocess data
    pass

def preprocess_module(queue_in, queue_res, args):
    while True:
        file = queue.get()
        # exit if getting terminating signal
        if file is None:
            break
        instance = data_preprocess(file, args)
        queue_res.put(instance)
    # put None into queue_res as terminating signal
    queue_res.put(None)


if __name__ == '__main__':
    nthread