PyTorch Dataloader freezes with num_workers > 0

Viewed 89

The following dataset class -> dataloader only works with num_workers = 0, and I'm not sure why. Other notebooks in the same environment do work with num_workers > 0. This has been bothering me for months!

Class that does not work: There is no error message, just runs indefinitely on next(iter(train_dl)), whereas with num_workers = 0 it takes 1sec.

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

In contrast, other notebooks using torchvision.datasets.ImageFolder(folder, transform) do work with num_workers > 0.

Any advice for how to make this compatible with async data loading, or other code feedback would be appreciated.

Python versoin 3.9.7 PyTorch version 1.10.1+cu113 Windows 11

0 Answers
Related