为什么 DataLoader 返回与 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

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

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

感谢您的时间。

Answer 1

它是由 collate_fn function 及其默认行为引起的。 它的主要目的是简化批次制备过程。 因此，您可以自定义您的批次制备过程，更新此 function。 As stated in documentationcollat e_fn , it automatically converts NumPy arrays and Python numerical values into PyTorch Tensors and it preserves the data structure. 所以它在你的情况下返回 [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

和 output 将是：

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

毕竟，我想再补充一点，你不应该只返回len function 中的 self.db.shape[0]。 在这种情况下，您的批量大小为 1，没关系； 但是，当批次大小发生变化时，它不会返回 #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

Answer 2

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

您返回从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)

此外，返回列表的长度是否应该与batch_size相同？

你为什么还要从DataLoader返回呢？ 只需从Dataset返回image ：

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

并且使用batch_size=12你会得到 output 的形状(12, 3, 60, 60) 。 您可以从此示例中获得形状，不要在Dataset中创建它，没有意义。