dataloader中的num_workers报错问题

PyTorch中自定义Dataloader的num_workers问题：解决训练loss为nan的异常现象

最新推荐文章于 2024-10-10 22:34:01 发布

原创最新推荐文章于 2024-10-10 22:34:01 发布 · 919 阅读

7 ·

CC 4.0 BY-SA版权

文章标签：

#python #pytorch

在使用PyTorch实现Yolov8时，自定义dataset和dataloader遇到num_workers>0时loss为nan的问题，原因是多线程环境下mosaic图像增强可能导致读取错误。通过添加线程锁解决了这一问题。

项目场景：

解决pytorch自定义dataloader中num_workers>0时训练loss出异常状况

问题描述

在复现yolov8时，自定义dataset和dataloader，在设置num_workers时，为0能正常训练，为其他值时loss为nan

class mosaic4:
    def __init__(self, dataset, p=0):
        self.dataset = dataset
        self.p = p

    def __call__(self, data):
        h, w = data['img'].shape[-2:]
        if random.random() < self.p:
                merge_boxes = []
                index = [random.randint(0, len(self.dataset) - 1) for _ in range(3)]
                mosaic = [self.dataset.load_img_and_label(item) for item in index]
              
               ...

           
            merge_boxes = torch.cat(merge_boxes,dim=0)
            boxes, cls = merge_boxes.split([4, 1] ,dim=-1)
            data = {'cls': cls, 'boxes': xyxy2xywh(boxes), 'img': new_img}

        return data

原因分析：

我的dataset中有mosaic图像增强，分析可能由于num_works优化了这一段代码使得图像读取出现错误

解决方案：

对mosaic这一段代码进行上锁处理加入：
if random.random() < self.p:
lock = threading.Lock()
with lock:

class mosaic4:
    def __init__(self, dataset, p=0):
        self.dataset = dataset
        self.p = p

    def __call__(self, data):
        h, w = data['img'].shape[-2:]
        if random.random() < self.p:
            lock = threading.Lock()
            with lock:
                merge_boxes = []
                index = [random.randint(0, len(self.dataset) - 1) for _ in range(3)]
                mosaic = [self.dataset.load_img_and_label(item) for item in index]
              
               ...

           
            merge_boxes = torch.cat(merge_boxes,dim=0)
            boxes, cls = merge_boxes.split([4, 1] ,dim=-1)
            data = {'cls': cls, 'boxes': xyxy2xywh(boxes), 'img': new_img}

        return data