PyTorch自定义数据集

数据传递机制

我们首先回顾识别手写数字的程序：

...
Dataset = torchvision.datasets.MNIST(root='./mnist/', train=True, transform=transform, download=True,)
dataloader = torch.utils.data.DataLoader(dataset=Dataset, batch_size=64, shuffle=True)
...
for epoch in range(EPOCH):
    for i, (image, label) in enumerate(dataloader):
        ...

从上面的程序，我们可以知道，在PyTorch中，数据传递机制是这样的：

创建Dataset
Dataset传递给DataLoader
DataLoader迭代产生训练数据提供给模型

总结这个数据传递机制就是，Dataset负责建立索引到样本的映射，DataLoader负责以特定的方式从数据集中迭代的产生一个个batch的样本集合。在enumerate过程中实际上是dataloader按照其参数sampler规定的策略调用了其dataset的getitem方法(下文中将介绍该方法)。关于Dataloder和Dataset的关系，具体可参考博客PyTorch中Dataset, DataLoader, Sampler的关系

在上面的识别手写数字的例子中，数据集是直接下载的，但如果我们自己收集了一些数据，存在电脑文件夹里，我们该如何把这些数据变为可以在PyTorch框架下进行神经网络训练的数据集呢，即如何自定义数据集呢？

自定义数据集

torch.utils.data.Dataset 是一个表示数据集的抽象类。任何自定义的数据集都需要继承这个类并覆写相关方法。所谓数据集，其实就是一个负责处理索引(index)到样本(sample)映射的一个类(class)。Pytorch提供两种数据集： Map式数据集 Iterable式数据集。这里我们只介绍前者。

一个Map式的数据集必须要重写getitem(self, index)、 len(self) 两个内建方法，用来表示从索引到样本的映射(Map)。这样一个数据集dataset，举个例子，当使用dataset[idx]命令时，可以在你的硬盘中读取数据集中第idx张图片以及其标签(如果有的话); len(dataset)则会返回这个数据集的容量。

自定义数据集类的范式大致是这样的：

class CustomDataset(torch.utils.data.Dataset):#需要继承torch.utils.data.Dataset
    def __init__(self):
        # TODO
        # 1. Initialize file path or list of file names.
        pass
    def __getitem__(self, index):
        # TODO
        # 1. Read one data from file (e.g. using numpy.fromfile, PIL.Image.open).
        # 2. Preprocess the data (e.g. torchvision.Transform).
        # 3. Return a data pair (e.g. image and label).
        #这里需要注意的是，第一步：read one data，是一个data point
        pass
    def __len__(self):
        # You should change 0 to the total size of your dataset.
        return 0

根据这个范式，我们举一个例子。

实例

从kaggle官网下载dogsVScats的数据集（百度网盘下载链接见文末），该数据集包含test1文件夹和train文件夹，train文件夹中包含12500张猫的图片和12500张狗的图片，图片的文件名中带序号：

cat.0.jpg
cat.1.jpg
cat.2.jpg
...
cat.12499.jpg
dog.0.jpg
dog.1.jpg
dog.2.jpg
...
dog.12499.jpg

我们把其中前10000张猫的图片和10000张狗的图片作为训练集，把后面的2500张猫的图片和2500张狗的图片作为验证集。猫的label记为0，狗的label记为1。因为图片大小不一，所以，我们需要对图像进行transform。

import matplotlib.pyplot as plt
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import os
image_transform = transforms.Compose([
    transforms.Resize(256),               # 把图片resize为256*256
    transforms.RandomCrop(224),           # 随机裁剪224*224
    transforms.RandomHorizontalFlip(),    # 水平翻转
    transforms.ToTensor(),                # 将图像转为Tensor
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])   # 标准化
])

class DogVsCatDataset(Dataset):   # 创建一个叫做DogVsCatDataset的Dataset，继承自父类torch.utils.data.Dataset
    def __init__(self, root_dir, train=True, transform=None):
        """
        Args:
            root_dir (string): Directory with all the images.
            transform (callable, optional): Optional transform to be applied on a sample.
        """
        self.root_dir = root_dir
        self.img_path = os.listdir(self.root_dir)
        if train:
            self.img_path = list(filter(lambda x: int(x.split('.')[1]) < 10000, self.img_path))    # 划分训练集和验证集
        else:
            self.img_path = list(filter(lambda x: int(x.split('.')[1]) >= 10000, self.img_path))
        self.transform = transform

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

    def __getitem__(self, idx):
        image = Image.open(os.path.join(self.root_dir, self.img_path[idx]))
        label = 0 if self.img_path[idx].split('.')[0] == 'cat' else 1        # label, 猫为0，狗为1
        if self.transform:
            image = self.transform(image)
        label = torch.from_numpy(np.array([label]))
        return image, label

我们来测试一下：

if __name__ == '__main__':
    catanddog_dataset = DogVsCatDataset(root_dir='/Users/wangpeng/Desktop/train',
                                        train=False,
                                        transform=image_transform)
    train_loader = DataLoader(catanddog_dataset, batch_size=8, shuffle=True, num_workers=4)   # num_workers=4表示用4个线程读取数据
    image, label = iter(train_loader).next()   # iter()函数把train_loader变为迭代器，然后调用迭代器的next()方法
    sample = image[0].squeeze()
    sample = sample.permute((1, 2, 0)).numpy()
    sample *= [0.229, 0.224, 0.225]
    sample += [0.485, 0.456, 0.406]
    sample = np.clip(sample, 0, 1)
    plt.imshow(sample)
    plt.show()
    print('Label is: {}'.format(label[0].numpy()))

运行结果：

PyTorch自定义数据集