PyTorch DataLoader的VSCode错误?

时间:2018-12-06 22:11:44

标签: visual-studio-code pytorch

以下代码示例在Python中可用,但在Linux中的VSCode中失败(在Windows中不是VSCode)。我想知道我的代码是否有问题,或者在Linux下VSCode是否有问题?

#Test of PyTorch DataLoader and Visual Studio Code
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]

if __name__ == '__main__':

    #Initialize DataLoader with above Dataset:
    dataloader = DataLoader(SimpleData(), batch_size=4, num_workers=1)

    print('Using DataLoader to show data in batches: ')
    for i, sample_batch in enumerate(dataloader):    #This fails in VSCode in Linux
        print('batch ', i, ':', sample_batch)

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

预期输出为:

Using DataLoader to show data in batches:
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 ---

但是在Linux中使用VSCode时,它在打印第一行后挂起。

1 个答案:

答案 0 :(得分:2)

您是否尝试过num_workers=0?

可能是VS Code无法在Linux上正确产生新进程。