DataLoader hanging/crashing

The DataLoader class is hanging (or crashing) in Windows but not in Linux with the following example:

#Demo of DataLoader crashing in Windows and with Visual Studio Code
import torch
from torch.utils.data import Dataset, DataLoader

class SimpleData(Dataset):
    """Very simple dataset"""
    def __init__(self):
        self.data = range(20)
    def __len__(self):
        return len(self.data)
    def __getitem__(self, idx):
        return self.data[idx]

#Create dataset
myDataSet = SimpleData()

#put dataset into DataLoader:
dataloader = DataLoader(myDataSet, batch_size=4, num_workers=1)

print('Using DataLoader to show data: ')
for i, sample_batched in enumerate(dataloader):
    print('batch ', i, ':', sample_batched)

print("--- Done ---")

When run in python in Liux I get, this outputs:

Using DataLoader to show data:
batch  0 : tensor([0, 1, 2, 3])
batch  1 : tensor([4, 5, 6, 7])
batch  2 : tensor([ 8,  9, 10, 11])
batch  3 : tensor([12, 13, 14, 15])
batch  4 : tensor([16, 17, 18, 19])
--- Done ---

Now when I insert an if name==‘main’: line, it works fine, like this:

#Demo of DataLoader working
import torch
from torch.utils.data import Dataset, DataLoader

class SimpleData(Dataset):
    """Very simple dataset"""
    def __init__(self):
        self.data = range(20)
    def __len__(self):
        return len(self.data)
    def __getitem__(self, idx):
        return self.data[idx]

#This crashes unless I include this line:
if __name__ == '__main__':

    #Create dataset
    myDataSet = SimpleData()

    #put dataset into DataLoader:
    dataloader = DataLoader(myDataSet, batch_size=4, num_workers=1)

    print('Using DataLoader to show data: ')
    for i, sample_batched in enumerate(dataloader):
        print('batch ', i, ':', sample_batched)

    print("--- Done ---")

Can anyone explain what is wrong with the top example? Why is it necessary to check for main? I am getting some other strange behavior with DataLoader, and so want to at least understand this.

Windows needs this if-clause protection, since it’s using spawn for multiprocessing.
If you don’t protect your code, the script will be called multiple times executing all module-level code, thus yielding an error.
Have a look at the Windows FAQ for further information.

1 Like

Ah, thank you for the pointer. I guess the top example is not really proper Python style and so I should avoid it, even though it “works” in linux.

Actually I take it back. Neither the top nor the bottom example will run successfully in Visual Studio Code (under Linux). Both work in straight python in linux, and the lower works on Windows in VSCode.

So is there something wrong with the second example? Or is there something wrong with VSCode in linux?

I’m not sure what might be the issue, since I’m not using VS Code for Python.
Do you get any kind of error message or does the run just hang?

It just hangs. I got some advice to set num_workers=0, and that worked! Seems like VSCode cannot properly spawn multiple instances in Linux.

1 Like