欢迎您访问程序员文章站本站旨在为大家提供分享程序员计算机编程知识!
您现在的位置是: 首页

pytorch dataloader和batch_size大小的理解

程序员文章站 2022-06-13 10:02:20
...

pytorch中dataloader的大小将根据batch_size的大小自动调整。如果训练数据集有1000个样本,并且batch_size的大小为10,则dataloader的长度就是100。

需要注意的是,如果dataset的大小并不能被batch_size整除,则dataloader中最后一个batch可能比实际的batch_size要小。例如,对于1001个样本,batch_size的大小是10,train_loader的长度len(train_loader)=101,最后一个batch将仅含一个样本。可以通过设置dataloader的drop_last属性为True来避免这种情况。

Pytorch的Dataset提供了得到dataset中单个样本的接口,即通过使用index下标来获得单个样本数据。Dataloader用来提供batch个样本供模型训练

class MyDataset(Dataset):
    def __init__(self, size):
        self.x = torch.randn(size, 1)
    
    def __getitem__(self, index):
        return self.x[index]

    def __len__(self):
        return len(self.x)

dataset = MyDataset(1001)

data_loader = DataLoader(dataset,
                         batch_size=10)

len(data_loader)

for batch_idx, data in enumerate(data_loader):
    print 'batch idx{}, batch len {}'.format(
        batch_idx, len(data))

data_loader = DataLoader(dataset,
                     batch_size=10,
                     drop_last=True)

len(data_loader)

for batch_idx, data in enumerate(data_loader):
    print 'batch idx{}, batch len {}'.format(
        batch_idx, len(data))

参考:https://discuss.pytorch.org/t/about-the-relation-between-batch-size-and-length-of-data-loader/10510/4