PyTorch Dataset、DataLoader长度

pytorch 可以直接对 Dataset 对象用 len() 求数据集大小，而 DataLoader 对象也可以用 len()，不过求得的是用这个 loader 在一个 epoch 能有几多 iteration，容易混淆。本文记录几种情况的对比。

from torch.utils.data import Dataset, DataLoader
class DS(Dataset):
    def __init__(self):
        self.data = list(range(17))
    def __len__(self):
        return len(self.data)
    def __getitem__(self, index):
        return self.data[index]
dset = DS()
loader = DataLoader(dset, batch_size=3)
loader2 = DataLoader(dset, batch_size=3, drop_last=True)
print(len(dset), len(loader.dataset), len(loader2.dataset)) # 17, 17, 17
print(len(loader)) # 6
print(len(loader2)) # 5