为什么 DataLoader return 列表的长度与 batch_size 不同

Question

我正在编写自定义数据加载器，而返回值让我感到困惑。

import torch
import torch.nn as nn
import numpy as np
import torch.utils.data as data_utils

class TestDataset:
    def __init__(self):
        self.db = np.random.randn(20, 3, 60, 60)

    def __getitem__(self, idx):
        img = self.db[idx]
        return img, img.shape[1:]

    def __len__(self):
        return self.db.shape[0]


if __name__ == '__main__':
    test_dataset = TestDataset()
    test_dataloader = data_utils.DataLoader(test_dataset,
                                       batch_size=1,
                                       num_workers=4,
                                       shuffle=False, \
                                       pin_memory=True
                                       )
    for i, (imgs, sizes) in enumerate(test_dataloader):
        print(imgs.size())  # torch.Size([1, 3, 60, 60])
        print(sizes)  # [tensor([60]), tensor([60])]
        break

为什么“调整”returns 长度为 2 的列表？我认为它应该是“torch.Size([1, 2])”，表示图像的高度和宽度(1 batch_size).

另外，返回列表的长度是否应该与batch_size一样？如果我想得到大小，我必须写“sizes = [sizes[0][0].item(), sizes[1][0].item()]”。这让我很困惑。

感谢您的宝贵时间。

Answer 1

Why "sizes" returns a list of length 2?

您 return 从 db 切片的单个元素切片 shape。这个代码片段应该更清楚：

import numpy as np

db = np.random.randn(20, 3, 60, 60)
img = db[0]
img.shape # (3, 60, 60)
img.shape[1:] # (60, 60)

Further more, should the length of the returned list be the same to batch_size?

为什么你甚至 return 来自 DataLoader 的？ return image 来自 Dataset:

def __getitem__(self, idx):
    return self.db[idx]

并且使用 batch_size=12 你会得到形状 (12, 3, 60, 60) 的输出。你可以从这个样本中得到形状，不要在 Dataset 中创建它，没有意义。

Answer 2

这是由 collate_fn 函数及其默认行为引起的。它的主要目的是简化批量准备过程。因此，您可以自定义您的批次准备过程并更新此功能。如文档 collate_fn 中所述，它会自动将 NumPy 数组和 Python 数值转换为 PyTorch 张量，并保留数据结构。所以在你的情况下 returns [tensor([60]), tensor([60])]。在许多情况下，您 return 将标签作为张量（而不是图像大小）的图像前馈到神经网络。我不知道为什么你在枚举时 return 图像大小，但是你可以通过添加自定义 collate_fn 来获得你需要的东西：

def collate_fn(data):
    imgs, lengths = data[0][0],data[0][1]    
    return torch.tensor(imgs), torch.tensor([lengths])

那么你应该将它设置为DataLoader的参数：

test_dataloader = DataLoader(test_dataset,
                                    batch_size=1,
                                    num_workers=4,
                                    shuffle=False, \
                                    pin_memory=True, collate_fn=collate_fn
                                    )

然后你可以循环为：

for i, (imgs, sizes) in enumerate(test_dataloader):
    print(imgs.size())
    print(sizes)  
    print(sizes.size())  
    break

输出为：

torch.Size([3, 60, 60])
tensor([[60, 60]])
torch.Size([1, 2])

毕竟，我想再补充一件事，你不应该只是 return self.db.shape[0] in len 函数。在这种情况下，您的批量大小为 1，没关系；但是，当批量大小发生变化时，它不会 return #batches 的真实值。您可以将 class 更新为：

class TestDataset:
    def __init__(self, batch_size=1):
        self.db = np.random.randn(20, 3, 60, 60)
        self._batch_size = batch_size
        
    def __getitem__(self, idx):
        img = self.db[idx]
        return img, img.shape[1:]

    def __len__(self):
        return self.db.shape[0]/self._batch_size

为什么 DataLoader return 列表的长度与 batch_size 不同

Why DataLoader return list that has a different length with batch_size

python

pytorch

dataloader