pytorch dataloader和batch_size大小的理解

YoJayC發表於2020-10-28

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

相關文章