时间戳

首页 > 解决方案 > 为什么 DataLoader 返回与 batch_size 长度不同的列表

python - 为什么 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, 3, 60, 60])
        print(sizes)  # [tensor([60]), tensor([60])]
        break

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

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

感谢您的时间。

标签： pythonpytorchdataloader

解决方案

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

推荐阅读