代码收藏家技术教程 2022-07-23

Pytorch中DataLoader和Dataset的基本用法

文章目录

DataLoader支持的两种数据集

Iterator格式的DataLoader

Python的Iterator格式数据简介

Pytorch使用DataLoader

使用自定义的IterableDataset

实战：自定义图片加载DataLoader

Map格式的DataLoader

自定义Map类型的Dataset

参考资料

DataLoader支持的两种数据集

Map格式：即key,value形式，例如 {0: ‘张三’, 1: ‘李四’}
Iterator格式：例如数组，迭代器等

Iterator格式的DataLoader

Python中，只要可以for循环的数据，都是Iterator格式的数据。

Python的Iterator格式数据简介

data = [0,1,2,3,4]

for item in data:
    print(item, end=' ')

0 1 2 3 4

上例子中，list数据类型是一个迭代器，for循环本质是每次调用了next函数。即其“效果”等价于下面的代码：

data = [0,1,2,3,4]
data_iter = iter(data) # 返回一个迭代器

item = next(data_iter, None) # 获取迭代器的下一个值
while item is not None:
    print(item, end=' ')
    item = next(data_iter, None)

0 1 2 3 4

Pytorch使用DataLoader

from torch.utils.data import DataLoader

data = [i for i in range(100)] # 定义数据集，需要是一个可迭代的对象

"""
定义dataloader，其接受三个重要的参数
- dataset: 数据集
- batch_size: 要将数据集切分为多少份
- shuffle: 是否对数据集进行随机排序
"""
dataloader = DataLoader(dataset=data, batch_size=6, shuffle=False) 

for i, item in enumerate(dataloader): # 迭代输出
    print(i, item)

0 tensor([0, 1, 2, 3, 4, 5])
1 tensor([ 6,  7,  8,  9, 10, 11])
... 省略
15 tensor([90, 91, 92, 93, 94, 95])
16 tensor([96, 97, 98, 99])

上面例子中，输入一个数据集0~99，通过dataloader将数据集分成100/6 =17份，每份6个数据，最后一份因为不满6个，所以只返回了4个。

使用自定义的IterableDataset

from torch.utils.data import DataLoader
from torch.utils.data import IterableDataset

class MyDataset(IterableDataset):
    
    def __init__(self):
        print('init...')
    
    def __iter__(self):
        print('iter...') # 获取迭代器
        self.n = 1
        return self
    
    def __next__(self):
        print('next...') # 获取下一个元素
        x = self.n
        self.n += 1
        
        if x >= 100: # 当x到100时停止
            raise StopIteration
        return x

dataloader = DataLoader(MyDataset(), batch_size=5)

for i, item in enumerate(dataloader):
    print(i, item)

init...
iter...
next... next... next... next... next... 0 tensor([1, 2, 3, 4, 5])
next... next... next... next... next... 1 tensor([ 6,  7,  8,  9, 10])
... 省略
next... next... next... next... next... 18 tensor([91, 92, 93, 94, 95])
next... next... next... next... next... 19 tensor([96, 97, 98, 99])

从上面的例子可以看出，可迭代对象在初始化中会调用一次__init__方法，在获取迭代器时会调用一次__iter__方法，之后在获取元素时，每获取一个元素都会调用一次__next__方法

实战：自定义图片加载DataLoader

任务：从data\faces文件夹中读取图片，并做一定处理，然后通过dataloader加载。

数据集：
链接：https://pan.baidu.com/s/1UrygNmmfzcWdjb29JHEpFg 
提取码：toyd

1.定义ImageDataset

import os
from torch.utils.data import IterableDataset
import torchvision
import torchvision.transforms as transforms

class ImageDataset(IterableDataset):
    def __init__(self, filepath):
        fnames = [filepath + '/' + filename for filename in os.listdir(filepath)] # 读取所有图片的文件路径
        self.i = -1 # 记录当前读取到的图片的下标
        self.compose = compose = [  # 图片的transform
                            transforms.ToPILImage(),
                            transforms.Resize((64, 64)),
                            transforms.ToTensor(),
                            transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
                        ]
    
    def __len__(self):
        return len(fnames) # 假设文件夹没有其他无关文件
    
    def __iter__(self):
        return self
    
    def __next__(self):
        self.i += 1
        if self.i >= len(self.fnames):
            raise StopIteration
        img = torchvision.io.read_image(fnames[self.i]) # 读取第i个图片
        transform = transforms.Compose(self.compose) # 对图片进行处理
        return transform(img) # 返回处理后的图片

2.实例化dataset和dataloader

dataset = ImageDataset('./data/faces')
print(next(iter(dataset)).shape)

dataloader = DataLoader(dataset, batch_size=16)
print(dataloader)

torch.Size([3, 64, 64])
<torch.utils.data.dataloader.DataLoader object at 0x000001C10B6048E0>

3.使用dataloader

import matplotlib.pyplot as plt

grid_img = torchvision.utils.make_grid(next(iter(dataloader)), nrow=4)
plt.figure(figsize=(10,10))
plt.imshow(grid_img.permute(1, 2, 0))
plt.show()

Map格式的DataLoader

dataset = {0: '张三', 1:'李四', 2:'王五', 3:'赵六'}

dataloader = DataLoader(dataset, batch_size=2)

for i, value in enumerate(dataloader):
    print(i, value)

0 ['张三', '李四']
1 ['王五', '赵六']

自定义Map类型的Dataset

自定义Map类型的Dataset只需要定义类，并继承 torch.utils.data.Dataset 方法即可，但要实现两个重要方法：__getitem__(self, index) 和 __len__(self)

个人比较推荐使用这种Dataset

例如：

from torch.utils.data import Dataset
from torch.utils.data import DataLoader

class CustomerDataSet(Dataset):

    def __init__(self):
        super(CustomerDataSet, self).__init__()
        self.data_dict = ['张三', '李四', '王五', '赵六']

    def __getitem__(self, index):
        return self.data_dict[index]

    def __len__(self):
        return len(self.data_dict)

其实这和上面的dict有异曲同工之妙，Dataloader会根据你dataset的大小，然后传一个index (0<=index<len(dataset)) 给getitem 方法，然后你返回该index对应的数据即可。例如：

from torch.utils.data import DataLoader

dataloader = DataLoader(CustomerDataSet(), batch_size=2, shuffle=True)
for i, value in enumerate(dataloader):
    print(i, value)

输出：

0 ['李四', '张三']
1 ['王五', '赵六']

在上面例子中，Dataloader的执行过程为：

调用 len(dataset) 方法，获取dataset的长度，这里为 4
然后生成 index list，即 [0,1,2,3]
因为传了shuffle=True，所以将index顺序打乱，结果为：[1,0,2,3]
然后按照顺序调用getitem方法，即：getitem(1)、getitem(0)、getitem(2)、getitem(3)
根据batch_size进行返回，第一次返回两个getitem(1)、getitem(0)，第二次返回getitem(2)、getitem(3)