Pytorch多进程Queue通信产生Segmentation fault (core dumped)——解决方案及代码规范

最新推荐文章于 2024-05-28 23:52:11 发布

进击De小火柴

最新推荐文章于 2024-05-28 23:52:11 发布

阅读量2.3k

点赞数 4

分类专栏：深度学习文章标签： pytorch 强化学习多进程 multiprocess queue

本文链接：https://blog.csdn.net/xzy5210123/article/details/107446352

版权

在Pytorch多进程强化学习项目中，作者遇到Queue通信导致的Segmentation fault (core dumped)问题。通过排查发现，错误源于未正确处理Queue中的Tensor数据，导致内存耗尽。解决方案是遵循官方文档建议，及时释放消费者进程中的内存，保持生产者进程运行，避免直接使用接收到的tensor。

摘要由CSDN通过智能技术生成

最近在做一个强化学习的项目，运用多进程分布训练时遇到了段错误的问题，这里记录下解决的过程思路和方案。

由于智能体与环境交互的过程涉及到了第三方的程序以及大量的文件读写操作，使得整个实验过程非常慢，为了解决交互部分的速度瓶颈，采用Ape-X（ Distributed Prioritized Experience Replay）的分布式训练思路，即多个actor负责与环境交互，得到的交互数据存储到公共replay memory中，一个leaner负责从memory中抽样训练更新网络。

由于Pytorch在多进程方面的封装较好，我采用torch.multiprocessing包来实现多进程，并通过其中的Queue队列来实现进程间通信，也就是actor将交互数据发送给learner。主要代码结构简化如下：

def actor(q):
    # 创建环境
    ...
    while True:
        # 获取交互数据 batch 类型为Tensor
        ...
        q.put(batch)

def learner(q)
    # 创建memory
    memory = Memory()
    ...
    while True:
        batch = q.get() # <--- *** 产生 SegFault的地方 ***
        memory.push(batch) 
        update_model()

if __name__ == '__main__':
    # 创建模型、优化器等
    model = DQN()
    model.share_memory()
    ...