pytorch dataloader num_workers参数设置导致训练阻塞

问题描述:

最近在用RFBnet (源码是pytorch的)训练RSNA的比赛数据,除了要修改一点代码支持RSNA的数据集外(打算后续再写个博客),发现在使用dataloader读取数据时,如果设置num_workers为0,也就是用主进程读取数据,模型训练程序运行正常。如果设置num_workers为其他任何一个大于0的整数,也就是使用子进程读取数据时,训练程序会卡住,卡在训练之前,GPU使用率和显存占用率都为0。

 

解决过程:

由于是多进程就,用print调试大法,定位到是一行opencv代码的问题,在dataloader子类的__getitem__方法里面调用了

image_array = cv2.cvtColor(image_array, cv2.COLOR_GRAY2BGR)

所有子进程的代码都卡在这里了。之前也有遇到过类似的问题,python多进程调用opencv的代码会出现问题。于是就用numpy的concate替代了这个方法

image_array = np.concatenate([image_array,image_array,image_array],axis=2)

绕过这个问题就可以正常训练了。

 

问题探索:

其实这个问题,我在另外一台机器上是没有遇到了,该机器是python3.6+opencv3.4.2

遇到问题的环境是docker环境,python3.5+opencv3.2

我感觉跟opencv的版本可能有关系,等后续训练完成之后把docker里面opencv的版本升级到3.4.2再试一下

你可能感兴趣的:(深度学习,pytorch,dataloader,num_workers,深度学习探索,图像处理)