为什么 DataLoader 返回与 batch_size 长度不同的列表答案

【问题标题】：Why DataLoader return list that has a different length with batch_size为什么 DataLoader 返回与 batch_size 长度不同的列表
【发布时间】：2021-06-01 18:42:28
【问题描述】：

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

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

为什么“sizes”返回一个长度为 2 的列表？我认为它应该是“torch.Size([1, 2])”，它表示图像的高度和宽度（1 batch_size）。

更进一步，返回列表的长度是否应该与batch_size相同？如果我想得到尺寸，我必须写“sizes = [sizes[0][0].item(), sizes[1][0].item()]”。这让我很困惑。

感谢您的宝贵时间。

【问题讨论】：

标签： python pytorch dataloader

【解决方案1】：

为什么“sizes”返回一个长度为 2 的列表？

您返回从 db 切片的单个元素的切片 shape。这段代码 sn -p 应该更清楚：

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)

更进一步，如果返回列表的长度与批量大小？

您为什么还要从DataLoader 退回它？只需从Dataset 返回image：

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

使用batch_size=12，您将获得形状为(12, 3, 60, 60) 的输出。你可以从这个示例中获得形状，不要在Dataset 中创建它，没有意义。

【讨论】：

【解决方案2】：

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

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

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

【讨论】：