Pytorch中DataLoader和Dataset的基本用法
文章目录
DataLoader支持的两种数据集
- Map格式：即key,value形式，例如 {0: '张三', 1: '李四'}
- Iterator格式：例如数组，迭代器等
Iterator格式的DataLoader
Python中，只要可以for循环的数据，都是Iterator格式的数据。
Python的Iterator格式数据简介
```python
data = [0,1,2,3,4]
for item in data:
print(item, end=' ')
```
``` 0 1 2 3 4 ```
上例子中，list数据类型是一个迭代器，for循环本质是每次调用了next函数。即其“效果”等价于下面的代码：
```python
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
```python
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
```python
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
```python
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
```python
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
```python 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
```python
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
例如：
```python
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对应的数据即可。例如：
```python 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)
上述的Dataloader执行过程只是为了方便大家理解，但具体源码是不是这样我也不清楚
参考资料
官方文档：https://pytorch.org/docs/stable/data.html
官方文档：https://pytorch.org/tutorials/beginner/basics/data_tutorial.html
