PyTorch Dataset、DataLoader长度

最新推荐文章于 2024-07-06 14:11:57 发布

HackerTom

最新推荐文章于 2024-07-06 14:11:57 发布

阅读量844

点赞数 2

分类专栏：环境机器学习文章标签： pytorch python

本文链接：https://blog.csdn.net/HackerTom/article/details/136736499

版权

机器学习同时被 2 个专栏收录

120 篇文章 16 订阅

订阅专栏

环境

88 篇文章 4 订阅

订阅专栏

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