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

如何解决为什么 DataLoader 返回与 batch_size 长度不同的列表

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

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,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()]”。这让我很困惑。

感谢您抽出宝贵时间。

解决方法

这是由 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)
        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

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

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

import numpy as np

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

此外，返回列表的长度应该与批量大小？

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

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

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