Pytorch 快速详解如何构建自己的Dataset完成数据预处理(附详细过程)
Pytorch(五) 使用DataSet和DataLoader数据加载
在这篇文章中我已经简单的介绍了Dataset和DataLoader的简单用法,但是大多数实际情况中数据集的存储都没有那么简单,所以写了本文来记录一下如何自定义DataSet
介绍
在实际的案例当中,如图像分类等任务来说,我们需要训练的数据集往往是存储在一个文件夹中的,而数据集的存储格式都是类似的
以蚂蚁和蜜蜂图片数据集 hymenoptera_data 来举例
一般的数据集都会分为两个文件夹
-
train训练集 -
val测试集
打开训练集之后的数据存储又分为两种情况
情况1
对于图像分类来说, 肯定需要一个 label和一个img
有些数据集喜欢把它们分开成两个文件夹
img文件夹 中存放的是图片
label文件夹中存放的是标签,通常以txt文件来存储,文件名和图片名相同,而文件的内容代表了图片的标签
情况2
对于一些简单的数据集来说,可能不会把label和img分开存放
比如情况1中提到的蚂蚁蜜蜂数据集
ants目录下的全是蚂蚁的图片
bees文件夹下全是蜜蜂的图片
这里的文件夹名就代表了图片的label
不过常用的情况 是把图片的label包含在了图片的命名当中
如下图
自定义Dataset
Dataset的主要作用就是提供一种方式来获取数据和其label
自定义的Dataset需要满足如下两个功能
- 如何获取每一个数据和其label
- 告诉我们数据集一共有多少个数据
导入库
from PIL import Image
from torch.utils.data import Dataset
import os如果没有下载相应的库就用 pip下载一下
TensorDataset回顾
在我之前的文章中提到,对于简单的数据可以用TensorDataset来包装
而通过for循环也可以遍历取出Dataset的中的数据和 label 它的原理就是内置了一个方法可以通过 index 来获取到相应的数据
下面就写个小案例来回顾一下TensorDataset的使用
构建数据x和标签y
x = torch.tensor(
[[1, 2, 3], [4, 5, 6], [7, 8, 9], [1, 2, 3], [4, 5, 6], [7, 8, 9], [1, 2, 3], [4, 5, 6], [7, 8, 9], [1, 2, 3],
[4, 5, 6], [7, 8, 9]])
y = torch.tensor([44, 55, 66, 44, 55, 66, 44, 55, 66, 44, 55, 66])`
from torch.utils.data import TensorDataset
train_dataset = TensorDataset(x, y)
print(train_dataset[0])
显然,如上图,train_dataset中数据的存放格式也是一个数据加一个标签的元组形式,并且可以通过 index来获取
当然在实际问题中也可以之间用for循环来遍历
Dataset自定义实现
那么对于我们的图像数据来说,要想达到遍历取数据和label的效果,需要我们自定义Dataset 在这里我们需要达到的目标是 通过 index 可以获取到一个 包含图像数据和标签的元组
并且要知道数据集的中长度
因此Dataset的初步模板就出来了, 如下
class MyDataset(Dataset):
def __init__(self):
pass
def __getitem__(self, index):
pass
def __len__(self):
pass其中 getitem的作用是通过index来返回一个包含数据和标签的元组len的作用是返回数据集的长度init很明显是类构造器
下面就来一步一步的实现它
获取图片并显示
PIL的Image可以帮助我们解析并查看一张图片
本案例的猫狗数据集我放在了D:\Source\Datasets\cat_and_dog 简单的读取一个图片来看看
from PIL import Image
img_path = r'D:\Source\Datasets\cat_and_dog\train\cat.10.jpg'
img = Image.open(img_path)
img.show()
通过 Image.open(img_path)方法读取到的是一个PIL.JpegImagePlugin.JpegImageFile对象,它包含了很多东西,我们就把它当做图像的数据
显然,读取图像需要图像的完整路径,那么思路很明显,我们可以把图像的路径存成一个列表,然后就可以通过index来获取列表中的值,(这里取出的是图像的完整路径),然后再通过Image把它解析成数据,取出其标签,返回
完成getitem方法
这一步其实不难,重点在理解其思路
下图就达到了返回img图像列表的操作,最后只需要把图像和train_path拼接起来就可以获取到图像的完整路径
完整代码
代码如下
# -*- coding: utf-8 -*-
# @Time : 2021/1/31 11:01
# @Author : Tong Tianyu
# @File : demo.py
from PIL import Image
from torch.utils.data import Dataset
import os
class MyDataset(Dataset):
def __init__(self, data_path):
self.data_path = data_path
self.img_list = os.listdir(self.data_path)
def __getitem__(self, index):
img_title = self.img_list[index]
img_label = img_title.split('.')[0]
img_path = os.path.join(self.data_path, img_title)
img = Image.open(img_path)
return img, img_label
def __len__(self):
return len(self.img_list)
train_path = r'D:\Source\Datasets\cat_and_dog\train'
train_dataset = MyDataset(train_path)效果如下, 完成目标
