PyTorch DataLoader 多进程调用OpenCV光流时遇到EOFerror和非法内存访问

环境

本文所描述的问题极可能与环境、版本等有关
win11 23H2
Nvidia Driver 555.85
i7 12700F + RTX 4070 Ti
Python 3.11.9
CUDA Runtime 12.1
CUDNN 8.9.7
PyTorch 2.3.0+cu121
OpenCV-contrib-cuda 4.9.0.80
mmengine 0.10.4
mmsegmentation 1.2.2

运行时任务

在openmm框架下实现分割任务，其中数据预处理阶段有涉及到CUDA的Opencv Optical Flow计算。因此任务中存在Opencv和PyTorch并行请求CUDA的情况。

由于OpenCV的CUDA加速的光流算法是不可pickle对象，因此采用mmengine.utils.ManagerMixin和multiprocessing.manager联合实现跨datalodaer workers的opencv光流代理中间件。

错误提示

错误于不可知时间点同时发生。
1.multiprocessing.manager进程间通信报错[WinError 109] 管道已结束
2.OpticalFlow.calc执行时提示EOFerror
3.（某种情况下代理进程有机会返回错误信息）opencv error: (-217:Gpu API call) an illegal memory access was encountered in function 'cv::cuda::Stream::waitForCompletion'

光流操作伪代码

from torch.utils.data import Dataset
class Dataset:
	...
	def OpticalFlow(self, imaga_series):
		OF = cv2.cuda.NvidiaOpticalFlow_2_0.create(...)
		flow = OF.calc(cv2.cuda.GpuMat, cv2.cuda.GpuMat, cv2.cuda.GpuMat)[0]
		flow = OF.convertToFloat(flow, cv2.cuda.GpuMat).download()
		return flow
	
	def __get_item__(self, idx):
		...
		flow = self.OpticalFlow(...)
		...

光流代理伪代码

from multiprocessing.managers import BaseManager
class OpticalFlowMultiProcessManager(BaseManager):
    pass

from mmengine.utils import ManagerMixin
class OpticalFlow_GlobalProxy(ManagerMixin):
    def __init__(self, name:str, **kwargs):
        super().__init__(name)
        ...
        OpticalFlowMultiProcessManager.register(OF_Type, eval(OF_Type))
        PythonMultiprocessManager = OpticalFlowMultiProcessManager()
        PythonMultiprocessManager.start()
        self.OpticalFlowGlobalService = getattr(PythonMultiprocessManager, OF_Type)(**kwargs)
        
    def __call__(self):
        return self.OpticalFlowGlobalService

Runtime

torch.utils.data.DataLoader设定num_worker>0 batchsize>1 （充要）

正常训练，即有一定概率抛出错误。

问题分析

综合各项报错，应当能够定位到是代理执行光流操作时产生了显存管理问题。或者是光流调用不当，导致opencv报c++的错误，但opencv的报错设计不是很好，有时候是调用不当的问题，但也会显示memory error之类的代码。

作者此处通过对比消融已经确定光流输入的参数均为合法，故怀疑是更加底层的问题导致的。

考虑到运行时复现此问题的充要条件包括num_worker>0 batchsize>1，所以非常自然地想到了多进程并行有关的内容。这两项设定都将导致数据加载时并行地对光流代理进行调用。

对光流操作伪代码的context、PID、ID等进行观察，发现光流操作保持在同一个PID进程中操作，同一时刻存在多个image_series的id，在不同代码位置运行着，与多线程的特征非常相似（具体不明，请自行查阅manager实现）。由于多线程会带来进程上下文切换，当延伸到CUDA层面时，可能引入非法内存指针问题。

说人话就是CUDA不太能保证正确处理并发的多个光流计算请求。40系显卡的光流加速器是新引入的，有点内存管理问题也就不足为奇了。

解决方法 - 使用进程锁限制光流计算

通过使用全局进程锁，实现在同一时刻只有一个“线程”在有效执行光流计算。通过观察，能够看到每个代理请求会被一次性执行完之后，才会有下一个代理请求进入并开始处理。

from multiprocessing import Lock
MP_LOCK = Lock()	# Global Lock

class OpticalFlowLocked(Dataset):
	def OpticalFlow(self, imaga_series):
		with MP_LOCK:
			OF = cv2.cuda.NvidiaOpticalFlow_2_0.create(...)
			flow = OF.calc(cv2.cuda.GpuMat, cv2.cuda.GpuMat, cv2.cuda.GpuMat)[0]
			flow = OF.convertToFloat(flow, cv2.cuda.GpuMat).download()
			return flow