【目标检测之数据集加载】利用DataLoader加载已预处理后的数据集【附代码】

在前一篇文章中，已经通过继承Dataset预处理自己的数据集 ，接下来就是使用pytorch提供的DataLoader函数加载数据集。

【目标检测之数据集预处理】继承Dataset定义自己的数据集【附代码】_z240626191s的博客-CSDN博客在深度学习训练中，除了设计有效的卷积神经网络框架外，更重要的是数据的处理。在训练之前需要对训练数据进行预处理。比如在目标检测网络训练中，首先需要划分训练集和测试集，然后对标签、边界框等进行处理后才能送入网络进行训练，本文章以VOC数据集格式为例，对数据集进行预处理后送入目标检测网络进行训练。【附代码】https://blog.csdn.net/z240626191s/article/details/123108750?spm=1001.2014.3001.5502

目录

 DataLoader函数

定义自己的collate_fn

---

 DataLoader函数

DataLoader是一个可以迭代的函数，用来加载并返回已处理好的数据集。其主要目的是对自定义后的Dataset数据集根据batch_size,是否随机打乱等封装成batch后进行训练。torch对于数据加载处理过程大致是这样的：

①创建自己的Dataset函数对数据预处理

②创建DataLoader对象

③循环访问这个DataLoader对象，返回image，Label【对于目标检测则包含bbox信息】送入模型进行训练

DataLoader主要参数：

dataset (Dataset):加载数据集，即自己定义后的Dataset

batch_size (int, optional):batch大小，默认为1

shuffle (bool, optional):如果为True表示每个迭代随机打乱

num_workers (int, optional):几个经常来处理数据

pin_memory (bool, optional)： 如果设置为True，将tensors拷贝到CUDA中的固定内存

collate_fn (callable, optional):将一个list的sample组成一个mini-batch的函数

前面的一些参数都比较好理解并且会设置，主要是最后一个collate_fn参数需要自己定义，如果不自己定义，默认返回的时候需要保证类型为tensor

定义自己的collate_fn

先附代码

def my_dataset_collate(batch):
    images = []
    bboxes = []
    for img, box in batch:
        images.append(img)
        bboxes.append(box)
    images = np.array(images)
    return images, bboxes

这里的batch即为通过Dataset返回的值，接下来演示一下【MyDatasets函数在我上篇有定义】

train_dataset = MyDatasets(lines[:num_train], (300, 300), True)  # train_dataset返回数据集和标签(且这个是可以迭代的)
train_data = DataLoader(train_dataset, batch_size, True, num_workers=4, pin_memory=True, collate_fn=my_dataset_collate)

我们先再看一下通过MyDatasets得到的是什么：

 In [1]: train_dataset
Out[1]: <__main__.MyDatasets at 0x27a33fe6c18>

然后继续打印一下

len(train_dataset)
Out[9]: 799     # 训练集的长度

-------------------------------------------------------------------------------------------------------------------------

# 打印一下第一个数据

In [10]: train_dataset[0]
Out[10]:
(array([[[ 67.,  86.,  96., ...,  49.,  64.,  26.],
         [ 64.,  81.,  95., ...,  35., 123.,  95.],
         [ 61.,  79.,  94., ...,  27., 124., 126.],
         ...,
         [ 36.,  39.,  42., ..., 130., 130., 124.],
         [ 37.,  38.,  37., ..., 130., 135., 132.],
         [ 36.,  37.,  38., ..., 129., 134., 132.]],

        [[ 59.,  74.,  77., ..., -58., -47., -89.],
·······

 array([[0.04666667, 0.52666664, 0.4       , 0.99333334, 0.        ],
        [0.08666667, 0.06333333, 0.32      , 0.5       , 0.        ],
        [0.62      , 0.54      , 0.8333333 , 0.99      , 0.        ],
        [0.56      , 0.        , 0.93      , 0.50333333, 0.        ]],
       dtype=float32))

可以看出在train_dataset[0]中，有两部分组成，最前面是图像信息，第二部分是bbox信息【包含类信息】,每个train_dataset[0,1,2···]都相当于batch=1的数据。

整个train_dataset都是有上述两部分组成。接下来我们送入DataLoader中看的collate_fn有什么变化 。

  for img,box in train_dataset:
           images.append(img)
           bboxes.append(box)

img是图像信息，box是边界框和类信息。并将这些放入列表中。

对于列表images

array([[[[  67.,   86.,   96., ...,   49.,   64.,   26.],
         [  64.,   81.,   95., ...,   35.,  123.,   95.],
         [  61.,   79.,   94., ...,   27.,  124.,  126.],
         ...,
         [  36.,   39.,   42., ...,  130.,  130.,  124.],
         [  37.,   38.,   37., ...,  130.,  135.,  132.],
         [  36.,   37.,   38., ...,  129.,  134.,  132.]],

        [[  59.,   74.,   77., ...,  -58.,  -47.,  -89.],
         [  55.,   68.,   74., ...,  -72.,   12.,  -19.],
         [  50.,   63.,   69., ...,  -79.,   14.,   14.],
         ...,
         [  71.,   74.,   77., ...,  110.,  110.,  104.],
         [  72.,   73.,   72., ...,  110.,  115.,  112.],
         [  71.,   72.,   73., ...,  109.,  114.,  112.]],

对于列表boxes

 [array([[0.04666667, 0.52666664, 0.4       , 0.99333334, 0.        ],
        [0.08666667, 0.06333333, 0.32      , 0.5       , 0.        ],
        [0.56      , 0.        , 0.93      , 0.50333333, 0.        ],
        [0.62      , 0.54      , 0.8333333 , 0.99      , 0.        ]],
       dtype=float32),
 array([[0.26333332, 0.19666667, 0.9633333 , 0.9766667 , 0.        ]],
       dtype=float32),

······················

在images列表中的每个元素的shape都是(3,300,300) 3指的3通道，我们也可以将其中一个通道进行显示：

最终我们可以返回一个images和bbox的列表。可以将这两个列表进行打包后再送入网络训练。