亲宝软件园·资讯

展开

Pytorch建模过程中的DataLoader与Dataset示例详解

奥辰 人气:0

处理数据样本的代码会因为处理过程繁杂而变得混乱且难以维护,在理想情况下,我们希望数据预处理过程代码与我们的模型训练代码分离,以获得更好的可读性和模块化,为此,PyTorch提供了torch.utils.data.DataLoader 和 torch.utils.data.Dataset两个类用于数据处理。其中torch.utils.data.DataLoader用于将数据集进行打包封装成一个可迭代对象,torch.utils.data.Dataset存储有一些常用的数据集示例以及相关标签。

同时PyTorch针对不同的专业领域,也提供有不同的模块,例如 TorchText(自然语言处理), TorchVision(计算机视觉), TorchAudio(音频),这些模块中也都包含一些真实数据集示例。例如TorchVision模块中提供了CIFAR, COCO, FashionMNIST 数据集。

1 定义数据集

pytorch中提供两种风格的数据集定义方式:

In [54]:

from torch.utils.data import Dataset
from torch.utils.data import DataLoader

In [56]:

dataset = {0: '张三', 1:'李四', 2:'王五', 3:'赵六', 4:'陈七'}
dataloader = DataLoader(dataset, batch_size=2)
for i, value in enumerate(dataloader):
    print(i, value)
0 ['张三', '李四']
1 ['王五', '赵六']
2 ['陈七']

In [57]:

from torch.utils.data import DataLoader
from torch.utils.data import IterableDataset

In [58]:

dataset = [i for i in range(10)]
dataloader = DataLoader(dataset=dataset, batch_size=3, shuffle=True) 
for i, item in enumerate(dataloader): # 迭代输出
    print(i, item)
0 tensor([3, 1, 2])
1 tensor([9, 7, 5])
2 tensor([0, 8, 4])
3 tensor([6])

如下所示,我们有一个蚂蚁蜜蜂图像分类数据集,目录结构如下所示,下面我们结合这个数据集,分别介绍如何使用这两个类定义真实数据集。

data
└── hymenoptera_data
    ├── train
    │   ├── ants
    │   │   ├── 0013035.jpg
    │   │   ……
    │   └── bees
    │       ├── 1092977343_cb42b38d62.jpg
    │       ……
    └── val
        ├── ants
        │   ├── 10308379_1b6c72e180.jpg
        │   ……
        └── bees
            ├── 1032546534_06907fe3b3.jpg
            ……

1.2 Dataset类

自定义一个Dataset类,继承torch.utils.data.Dataset,且必须实现下面三个方法:

In [211]:

import os
import pandas as pd
from PIL import Image
from torchvision.transforms import ToTensor, Lambda
from torchvision import transforms
import torchvision
class AntBeeDataset(Dataset):
    # 把图片所在的文件夹路径分成两个部分,一部分是根目录,一部分是标签目录,这是因为标签目录的名称我们需要用到
    def __init__(self, root_dir, transform=None, target_transform=None):
        """
        root_dir:存放数据的根目录,即:data/hymenoptera_data
        transform: 对图像数据进行处理,例如,将图片转换为Tensor、图片的维度可能不一致需要进行resize
        target_transform:对标签数据进行处理,例如,将文本标签转换为数值
        """
        self.root_dir = root_dir
        self.transform = transform
        self.target_transform = target_transform
        
        # 获取文件夹下所有图片的名称和对应的标签
        self.img_lst = []
        for label in ['ants', 'bees']:
            path = os.path.join(root_dir, label)
            for img_name in os.listdir(path):
                self.img_lst.append((os.path.join(root_dir, label, img_name), label))
        
    def __getitem__(self, idx):
        img_path, label = self.img_lst[idx]
        img = Image.open(img_path).convert('RGB')
        
        if self.transform:
            img = self.transform(img)
        if self.target_transform:
            label = self.target_transform(label)
        # 这个地方要注意,我们在计算loss的时候用交叉熵nn.CrossEntropyLoss()
        # 交叉熵的输入有两个,一个是模型的输出outputs,一个是标签targets,注意targets是一维tensor
        # 例如batchsize如果是2,ants的targets的应该[0,0],而不是[[0][0]]
        # 因此label要返回0,而不是[0]
        return img, label

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

In [310]:

train_transform = transforms.Compose([
    
    transforms.RandomResizedCrop(224),  # 将给定图像随机裁剪为不同的大小和宽高比,然后缩放所裁剪得到的图像为制定的大小
    transforms.RandomHorizontalFlip(),  # 以给定的概率随机水平旋转给定的PIL的图像,默认为0.5
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

# 验证集并不需要做与训练集相同的处理,所有,通常使用更加简单的transformer
val_transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

# 根据标签目录的名称来确定图片是哪一类,如果是"ants",标签设置为0,如果是"bees",标签设置为1
target_transform = transforms.Lambda(lambda y: 0 if y == "ants" else 1)

In [311]:

train_dataset = AntBeeDataset('data/hymenoptera_data/train', transform=train_transform, target_transform=target_transform)
val_dataset = AntBeeDataset('data/hymenoptera_data/val', transform=val_transform, target_transform=target_transform)

1.2 Dataset数据集常用操作

1. 查看数据集大小:

In [221]:

len(train_dataset), len(val_dataset)

Out[221]:

(245, 153)

2. 合并数据集

In [222]:

dataset = train_dataset + val_dataset

In [223]:

len(dataset)

Out[223]:

398

3. 划分训练集、测试集

In [224]:

from torch.utils.data import random_split
# random_split 不能直接使用百分比划分,必须指定具体数字
train_size = int( len(dataset) * 0.8)
test_size = len(dataset) - train_size

In [225]:

train_dataset, val_dataset = random_split(dataset, [train_size, test_size])

In [226]:

len(train_dataset), len(val_dataset)

Out[226]:

(318, 80)

1.3 IterableDataset类

使用迭代器风格时,必须继承IterableDataset类,且实现下面两个方法:

虽然只需要实现这两个方法,但是通常还需要在迭代过程中对数据进行处理。IterableDataset类实现自定义数据集,本质就是创建一个数据集类,且实现__iter__返回一个迭代器。一下提供两种方法通过IterableDataset类自定义数据集:

方法一:

In [289]:

class AntBeeIterableDataset(IterableDataset):
    # 把图片所在的文件夹路径分成两个部分,一部分是根目录,一部分是标签目录,这是因为标签目录的名称我们需要用到
    def __init__(self, root_dir, transform=None, target_transform=None):
        """
        root_dir:存放数据的根目录,即:data/hymenoptera_data
        transform: 对图像数据进行处理,例如,将图片转换为Tensor、图片的维度可能不一致需要进行resize
        target_transform:对标签数据进行处理,例如,将文本标签转换为数值
        """
        self.root_dir = root_dir
        self.transform = transform
        self.target_transform = target_transform
        
        # 获取文件夹下所有图片的名称和对应的标签
        self.img_lst = []
        for label in ['ants', 'bees']:
            path = os.path.join(root_dir, label)
            for img_name in os.listdir(path):
                self.img_lst.append((os.path.join(root_dir, label, img_name), label))
                
    def __iter__(self):
        for img_path, label in self.img_lst:
            img = Image.open(img_path).convert('RGB')
            if self.transform:
                img = self.transform(img)
            if self.target_transform:
                label = self.target_transform(label)
            yield img, label

方法二:

In [285]:

class AntBeeIterableDataset(IterableDataset):
    # 把图片所在的文件夹路径分成两个部分,一部分是根目录,一部分是标签目录,这是因为标签目录的名称我们需要用到
    def __init__(self, root_dir, transform=None, target_transform=None):
        """
        root_dir:存放数据的根目录,即:data/hymenoptera_data
        transform: 对图像数据进行处理,例如,将图片转换为Tensor、图片的维度可能不一致需要进行resize
        target_transform:对标签数据进行处理,例如,将文本标签转换为数值
        """
        self.root_dir = root_dir
        self.transform = transform
        self.target_transform = target_transform
        
        # 获取文件夹下所有图片的名称和对应的标签
        self.img_lst = []
        for label in ['ants', 'bees']:
            path = os.path.join(root_dir, label)
            for img_name in os.listdir(path):
                self.img_lst.append((os.path.join(root_dir, label, img_name), label))
        self.index = 0
                
    def __iter__(self):
        return self
    
    def __next__(self):
        try:
            img_path, label = self.img_lst[self.index]
            self.index += 1
            img = Image.open(img_path).convert('RGB')
            if self.transform:
                img = self.transform(img)
            if self.target_transform:
                label = self.target_transform(label)
            return img, label
        except IndexError:
            raise StopIteration()

In [290]:

train_dataset = AntBeeIterableDataset('data/hymenoptera_data/train', transform=train_transform, target_transform=target_transform)
val_dataset = AntBeeIterableDataset('data/hymenoptera_data/val', transform=val_transform, target_transform=target_transform)

在处理大数据集时,IterableDataset会比Dataset更有优势,例如数据存储在文件或者数据库中,只需要在自定义的IterableDataset之类中获取文件操作句柄或者数据库连接和游标惊喜迭代,每次只返回一条数据即可。我们把上文中蚂蚁蜜蜂数据集的所有图片、标签这里后写入hymenoptera_data.txt中,内容如下所示,假设有数亿行,那么,就不能直接将数据加载到内存了:

data/hymenoptera_data/train/ants/2288481644_83ff7e4572.jpg, ants
data/hymenoptera_data/train/ants/2278278459_6b99605e50.jpg, ants
data/hymenoptera_data/train/ants/543417860_b14237f569.jpg, ants
...
...

可以参考一下方式定义IterableDataset子类:

In [299]:

class AntBeeIterableDataset(IterableDataset):
    # 把图片所在的文件夹路径分成两个部分,一部分是根目录,一部分是标签目录,这是因为标签目录的名称我们需要用到
    def __init__(self, filepath, transform=None, target_transform=None):
        """
        filepath:hymenoptera_data.txt完整路径
        transform: 对图像数据进行处理,例如,将图片转换为Tensor、图片的维度可能不一致需要进行resize
        target_transform:对标签数据进行处理,例如,将文本标签转换为数值
        """
        self.filepath = filepath
        self.transform = transform
        self.target_transform = target_transform

                
    def __iter__(self):
        with open(self.filepath, 'r') as f:
            for line in f:
                img_path, label = line.replace('\n', '').split(', ')
                img = Image.open(img_path).convert('RGB')
                if self.transform:
                    img = self.transform(img)
                if self.target_transform:
                    label = self.target_transform(label)
                yield img, label

In [307]:

train_dataset = AntBeeIterableDataset('hymenoptera_data.txt', transform=train_transform, target_transform=target_transform)

注意,IterableDataset方法在处理大数据集时确实比Dataset更有优势,但是,IterableDataset在迭代过程中,样本输出顺序是固定的,在使用DataLoader进行加载时,无法使用shuffle进行打乱,同时,因为在IterableDataset中并未强制限定必须实现__len__()方法(很多时候确实也没法获取数据总量),不能通过len()方法获取数据总量。

2 DataLoad

DataLoader的功能是构建可迭代的数据装载器,在训练的时候,每一个for循环,每一次Iteration,就是从DataLoader中获取一个batch_size大小的数据,节省内存的同时,它还可以实现多进程、数据打乱等处理。我们通过一张图来了解DataLoader数据读取机制:

首先,在for循环中使用了DataLoader,进入DataLoader后,首先根据是否使用多进程DataLoaderIter,做出判断之后单线程还是多线程,接着使用Sampler得索引Index,然后将索引给到DatasetFetcher,在这里面调用Dataset,根据索引,通过getitem得到实际的数据和标签,得到一个batch size大小的数据后,通过collate_fn函数整理成一个Batch Data的形式输入到模型去训练。

在pytorch建模的数据处理、加载流程中,DataLoader应该算是最核心的一步操作DataLoader有很多参数,这里我们列出常用的几个:

我们重点说说多进程模式下使用DataLoader,在多进程模式下,每次 DataLoader 创建 iterator 时(遍历DataLoader时,例如,当调用时enumerate(dataloader)),都会创建 num_workers 工作进程。dataset, collate_fn, worker_init_fn 都会被传到每个worker中,每个worker都用独立的进程。

对于映射风格的数据集,即Dataset子类,主线程会用Sampler(采样器)产生indice,并将它们送到进程里。因此,shuffle是在主线程做的

对于迭代器风格的数据集,即IterableDataset子类,因为每个进程都有相同的data复制样本,并在各个进程里进行不同的操作,以防止每个进程输出的数据是重复的,所以一般用 torch.utils.data.get_worker_info() 来进行辅助处理。

这里,torch.utils.data.get_worker_info() 返回worker进程的一些信息(id, dataset, num_workers, seed),如果在主线程跑的话返回None

注意,通常不建议在多进程加载中返回CUDA张量,因为在使用CUDA和在多处理中共享CUDA张量时存在许多微妙之处(文档中提出:只要接收过程保留张量的副本,就需要发送过程来保留原始张量)。建议采用 pin_memory=True ,以将数据快速传输到支持CUDA的GPU。简而言之,不建议在使用多线程的情况下返回CUDA的tensor。

In [313]:

dataload = DataLoader(train_dataset, batch_size=2)

In [315]:

img, label = next(iter(dataload))

In [316]:

img.shape, label

Out[316]:

(torch.Size([2, 3, 224, 224]), tensor([0, 0]))

加载全部内容

相关教程
猜你喜欢
用户评论