多线程使用OpenCV出现死锁

问题复现

在使用Pytorch的dataloader时,有一个参数num_workers,设置这个参数为大于0的整数,可以开启多线程加载数据。多线程加载数据可以预先加载下几批数据进内存,目的是为了避免GPU运算完一个batch的数据后等待CPU计算浪费GPU算力。开启后GPU运算一批batch后会从内存中读取下一批已经完成预处理的数据,这样会使得GPU始终保持运算状态,加快整个训练的过程。

但是实际上,如果预处理使用了OpenCV模块,由于OpenCV默认会开启多线程进行处理,在Python没有开启多线程的时候不会出现问题。一旦使用num_workers>0时,Python也开启了多线程,会导致OpenCV的多线程互相冲突,形成死锁,会使程序卡死。

解决方案

手动设置OpenCV只使用单个线程,保证num_workers之间不会形成依赖,就可以打破死锁的形成条件,就能解决这个问题。

在主程序中加入如下代码即可。

import cv2
cv2.setNumThreads(0)