PyTorch 数据加载器因 num_workers > 0 冻结

Question

以下数据集类 -> dataloader 仅适用于 num_workers = 0，我不确定为什么。同一环境中的其他笔记本确实可以在 num_workers > 0 的情况下工作。这已经困扰我好几个月了！

不起作用的类：没有错误消息，只是在 next(iter(train_dl)) 上无限期地运行，而 num_workers = 0 则需要 1 秒。

class SegmentationDataSet(data.Dataset):
def __init__(self, fnames, rle_df=None, path=train_val_dir):
    
    self.fnames = fnames
    self.rle_df = rle_df
    self.path = path

def __len__(self):
    return len(self.fnames)

def __getitem__(self, index:int):           
    img_id = self.fnames[index]
    mask = None
    im = torchvision.io.read_image(self.path + img_id).float()
    
    if self.rle_df is not None:
        rle = self.rle_df.loc[self.rle_df['id']==img_id]['rle']

        if not pd.isnull(rle).values[0]:
            rle = rle.values[0]
            mask = rle2mask(rle, [1024,1024])
            mask = torch.from_numpy(np.expand_dims(mask,0))
        else:
            mask = torch.zeros([1,1024,1024])
            
    return self.transform(im, mask)
    
def transform(self, im, mask):
    im = im / 255
    im =  torchvision.transforms.Resize((512,512))(im)
    if mask is not None:
        mask = torchvision.transforms.Resize((512,512))(mask)
        return im, mask
    else:
        return im

相比之下，其他使用 torchvision.datasets.ImageFolder(folder, transform) 的笔记本可以在 num_workers > 0 的情况下工作。

任何有关如何使其与异步数据加载兼容的建议或其他代码反馈将不胜感激。

Python版本3.9.7 PyTorch版本1.10.1+cu113 Windows 11

Answer 1

遇到类似的问题后，看起来该问题与 Jupyter Notebooks 在 Windows 上处理多处理的方式有关（来源）。要解决此问题，请尝试在 Python 脚本中运行代码，或者如果您需要在笔记本中运行代码，请从 WSL 运行笔记本