Pytorch自定义Dataset
- 1. 自定义加载数据
- 1.1. 第一种 Dataset class
- 1.2. 第二种 torchvision
1. 自定义加载数据
在学习Pytorch的教程时，加载数据许多时候都是直接调用torchvision.datasets里面集成的数据集，直接在线下载，然后使用torch.utils.data.DataLoader进行加载。
那么，我们怎么使用我们自己的数据集，然后用DataLoader进行加载呢？
常见的两种形式的导入：
- 一种是整个数据集都在一个文件下，内部再另附一个label文件，说明每个文件的状态。这种存放数据的方式可能更时候在非分类问题上得到应用。
- 一种则是更适合在分类问题上，即把不同种类的数据分为不同的文件夹存放起来。这样，我们可以从文件夹或文件名得到label。
我们以猫狗数据集为例，进行自定义加载数据。
猫狗数据集里面有两个文件夹，分别是test和train。
其中train文件夹下的图片，命名方式为：cat.0.jpg或dog.0.jpg。我们可以从文件名中提取出来作为我们的图片标签。
1.1. 第一种 Dataset class
这种方法是官方导航介绍的。torch.utils.data.Dataset是一个抽象类，用户想要加载自定义的数据只需要继承这个类，并且覆写其中的两个方法即可：
-
__len__:实现len(dataset)返回整个数据集的大小。 -
__getitem__用来获取一些索引的数据，使dataset[i]返回数据集中第i个样本。 - 不覆写这两个方法会直接返回错误。
def __getitem__(self, index):
raise NotImplementedError
def __len__(self):
raise NotImplementedError建立的自定义类如下：
#导入相关模块
from torch.utils.data import DataLoader,Dataset
from skimage import io,transform
import matplotlib.pyplot as plt
import os
import torch
from torchvision import transforms
import numpy as np
class AnimalData(Dataset): #继承Dataset
def __init__(self, root_dir, transform=None): #__init__是初始化该类的一些基础参数
self.root_dir = root_dir #文件目录
self.transform = transform #变换
self.images = os.listdir(self.root_dir)#目录里的所有文件
def __len__(self):#返回整个数据集的大小
return len(self.images)
def __getitem__(self,index):#根据索引index返回dataset[index]
image_index = self.images[index]#根据索引index获取该图片
img_path = os.path.join(self.root_dir, image_index)#获取索引为index的图片的路径名
img = io.imread(img_path)# 读取该图片
label = img_path.split('\\')[-1].split('.')[0]# 根据该图片的路径名获取该图片的label，具体根据路径名进行分割。我这里是"E:\\Python Project\\Pytorch\\dogs-vs-cats\\train\\cat.0.jpg"，所以先用"\\"分割，选取最后一个为['cat.0.jpg']，然后使用"."分割，选取[cat]作为该图片的标签
sample = {'image':img,'label':label}#根据图片和标签创建字典
if self.transform:
sample = self.transform(sample)#对样本进行变换
return sample #返回该样本设置好数据类之后，我们就可以将其用torch.utils.data.DataLoader加载，并访问它。
if __name__=='__main__':
data = AnimalData('E:/Python Project/PyTorch/dogs-vs-cats/train',transform=None)#初始化类，设置数据集所在路径以及变换
dataloader = DataLoader(data,batch_size=128,shuffle=True)#使用DataLoader加载数据
for i_batch,batch_data in enumerate(dataloader):
print(i_batch)#打印batch编号
print(batch_data['image'].size())#打印该batch里面图片的大小
print(batch_data['label'])#打印该batch里面图片的标签输出如下：
0
torch.Size([128, 3, 224, 224])
['dog', 'cat', 'dog', 'dog', 'cat', 'dog', 'cat', 'cat', 'dog', 'dog', 'dog', 'dog', 'cat', 'cat', 'dog', 'dog', 'cat', 'dog', 'cat', 'cat', 'cat', 'cat', 'dog', 'cat', 'cat', 'cat', 'cat', 'cat', 'dog', 'dog', 'cat', 'cat', 'dog', 'cat', 'cat', 'cat', 'cat', 'dog', 'dog', 'dog', 'cat', 'dog', 'cat', 'cat', 'cat', 'dog', 'cat', 'dog', 'dog', 'dog', 'dog', 'dog', 'cat', 'cat', 'dog', 'cat', 'cat', 'cat', 'cat', 'dog', 'dog', 'cat', 'dog', 'cat', 'cat', 'cat', 'cat', 'dog', 'cat', 'dog', 'dog', 'cat', 'dog', 'dog', 'cat', 'cat', 'dog', 'dog', 'dog', 'dog', 'dog', 'cat', 'cat', 'dog', 'dog', 'cat', 'dog', 'cat', 'dog', 'cat', 'cat', 'dog', 'cat', 'dog', 'cat', 'dog', 'cat', 'dog', 'cat', 'cat', 'cat', 'dog', 'cat', 'cat', 'dog', 'dog', 'dog', 'cat', 'cat', 'cat', 'cat', 'dog', 'cat', 'dog', 'dog', 'dog', 'dog', 'cat', 'dog', 'dog', 'dog', 'dog', 'cat', 'cat', 'dog', 'cat', 'dog', 'cat']1.2. 第二种 torchvision
pytorch几乎将上述所有工作都封装起来供我们使用，其中一个工具就是torchvision.datasets.ImageFolder,用于加载用户自定义的数据，要求我们的数据要有如下结构：
将图片按类别分文件夹存放。
root/ants/xxx.png
root/ants/xxy.jpeg
root/ants/xxz.png
.
.
.
root/bees/123.jpg
root/bees/nsdf3.png
root/bees/asd932_.png
torchvision.transforms中也封装了各种各样的数据处理的工具，如Resize, ToTensor等等功能供我们使用。
from torchvision import transforms,utils
from torchvision import datasets
import torch
import matplotlib.pyplot as plt
import torch.utils.data
train_data = datasets.ImageFolder(r'E:\Python Project\PyTorch\data\hotdog\train',transform=transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor()
]))
print(train_data.classes)#获取标签
train_loader = torch.utils.data.DataLoader(train_data,batch_size=4,shuffle=True)
print(len(train_loader))
for i_batch, img in enumerate(train_loader):
if i_batch == 0:
print(img[1]) #标签转化为编码
fig = plt.figure()
grid = utils.make_grid(img[0])
plt.imshow(grid.numpy().transpose((1, 2, 0)))
plt.show()
break输出结果如下:
classes= ['hotdog', 'not-hotdog']
class_to_idx= {'hotdog': 0, 'not-hotdog': 1}
500
img[1]= tensor([0, 1, 0, 1])
本文章为转载内容，我们尊重原作者对文章享有的著作权。如有内容错误或侵权问题，欢迎原作者联系我们进行内容更正或删除文章。
