Dataset与Dataloader

在pytorch中，Dataset和Dataloader是用来加载数据的两个类

数据集

采用蚂蚁蜜蜂数据集，数据集的目录结构如下

hymenoptera_data/
|-- train/
|   |-- ants/
|   |   |-- 0013035.jpg
|   |   |-- 24335309_c5ea483bb8.jpg
|   |   |-- ... ...
|   |-- bees/
|   |   |-- 16838648_415acd9e3f.jpg
|   |   |-- ... ...
|-- val/
|   |-- ants/
|   |   |-- 8124241_36b290d372.jpg
|   |   |-- ... ...
|   |-- bees/
|   |   |-- 26589803_5ba7000313.jpg
|   |   |-- ... ...

Dataset

首先来看一下Dataset的函数文档

from torch.utils.data import Dataset
help(Dataset)
---------------------------------
class Dataset(typing.Generic)
 |  Dataset(*args, **kwds)
 |  
 |  An abstract class representing a :class:`Dataset`.
 |  
 |  All datasets that represent a map from keys to data samples should subclass
 |  it. All subclasses should overwrite :meth:`__getitem__`, supporting fetching a
 |  data sample for a given key. Subclasses could also optionally overwrite
 |  :meth:`__len__`, which is expected to return the size of the dataset by many
 |  :class:`~torch.utils.data.Sampler` implementations and the default options
 |  of :class:`~torch.utils.data.DataLoader`.

可以看到，Dataset是一个抽象类（接口），实现该接口的子类必须实现__getitem__和__len__方法

接下来通过Dataset实现对蚂蚁蜜蜂数据集的加载

from torch.utils.data import Dataset
from PIL import Image
import os

class MyData(Dataset):
    def __init__(self, root, is_ants=True, is_train=True):
        if is_ants:
            self.label = 'ants'
        else:
            self.label = 'bees'
        if is_train:
            tv = 'train'
        else:
            tv = 'val'
        self.path = os.path.join(root, tv, self.label) 
        self.img_list = os.listdir(self.path)
        
    def __getitem__(self, idx):
        img_name = self.img_list[idx]
        img_path = os.path.join(self.path, img_name)
        img = Image.open(img_path)
        return img, self.label  # 返回第idx张图片和该图片的标签
    
    def __len__(self):
        return len(self.img_list)

使用MyData对数据集进行加载和读取

val_ants = MyData(r'D:\dataset\hymenoptera_data', True, False)  # 加载验证集中的蚂蚁数据集
print("共%d个样本" % len(val_ants))
print(val_ants[0])  # 读取第10张图片及其标签
val_bees = MyData(r'D:\dataset\hymenoptera_data', False, False)  # 加载验证集中的蜜蜂数据集
print("共%d个样本" % len(val_bees))
print(val_bees[0])  # 读取第10张图片及其标签
---------------------------------
共70个样本
(<PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=500x375 at 0x1D58C7C0550>, 'ants')
共83个样本
(<PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=366x500 at 0x1D58C7C08E0>, 'bees')

由于Dataset中已经定义了__add__方法，因此我们可以直接通过加法将两个数据集合并

val_data = val_ants + val_bees
print("共%d个样本" % len(val_data))
print(val_data[len(val_ants)-1][1], val_data[len(val_ants)][1])
---------------------------------
共153个样本
ants bees