数据集定义与加载#

作者: MinYao Ni

深度学习模型的训练与评估均需要数据,数据根据实际模型需求可以是图片、文本、语音等类型。在训练前,数据会经过一系列处理,转换成模型可计算的类型和分布,然后批量送入模型;这需要:
1.定义数据集:将原始数据和标签映射到Dataset,方便后续通过索引读取数据,在读取的同时也可以进行数据预处理与数据增强等操作。
2.定义数据迭代器:将定义好的数据集打乱顺序、分批(也可以进行mixup等更多操作),方便训练和评估时读取。

定义数据集#

加载内置数据集#

框架在 neurai.datasets 下内置了一部分经典数据集以供开发者调用,如MNIST系列、CIFAR10等数据集。
以MNIST数据集为例,加载内置数据集流程如下:
from neurai import datasets

train_data = datasets.MNIST(root="./", train=True, download=True)
test_data = datasets.MNIST(root="./", train=False, download=True)
print('train images: ',len(train_data),', test images: ',len(test_data))
train images:  60000 , test images:  10000
参数 download 值为 True 时,如果目录“./”中无数据集,将自动下载数据至“./”;
MNIST数据集已经划分好了训练集和测试集, 通过 train 参数区分,参数为 True 时加载训练数据, False 时加载测试数据。
完成数据集初始化之后,可以使用下面的代码直接对数据集进行迭代读取。
import cv2

for data in train_data:
  image, label = data
  print('shape of image: ',image.shape)
  cv2.imwrite(str(label)+".png", image)
  break
shape of image:  (28, 28, 1)
../_images/mnist-5.png

自定义数据集#

在实际的开发场景中,很多需要使用自己收集的数据来进行训练和验证,此时可以通过继承 datasets.Dataset 来自定义数据集,如果是图像数据集也可以继承 datasets.VisionDataset
可基于以上两个基类构建子类,并实现下面3个函数:
__init__ : 完成数据集初始化操作,将磁盘中的样本文件路径映射到一个列表中。
__getitem__ :完成根据索引获取数据和标签,并返回该数据对(样本数据、对应的标签)。
__len__ :返回数据集的样本总数。
下面不使用MNIST预置类,自定义MNIST数据集作为示例。
import os
import cv2
from neurai import datasets
from neurai.datasets.transforms import Normalize

class MyDataset(datasets.Dataset):
  """
  步骤一:继承 datasets.Dataset 类
  """
  def __init__(self, root, label_path, transform=None):
    """
    步骤二:实现 __init__ 函数,初始化数据集,将样本和标签映射到列表中
    """
    super().__init__()
    self.data_list = []
    with open(label_path,encoding='utf-8') as f:
      for line in f.readlines():
          image_path, label = line.strip().split('\t')
          image_path = os.path.join(root, image_path)
          self.data_list.append([image_path, label])
    # 传入定义好的数据处理方法,作为自定义数据集类的一个属性
    self.transform = transform

  def __getitem__(self, index):
    """
    步骤三:实现 __getitem__ 函数,定义指定 index 时如何获取数据,并返回单条数据(样本数据、对应的标签)
    """
    # 根据索引,从列表中取出一个图像
    image_path, label = self.data_list[index]
    # 读取灰度图
    image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    # 将图像数据格式转换为 float32
    image = image.astype('float32')[...,None]
    # 应用数据处理方法到图像上
    if self.transform is not None:
        image = self.transform(image)
    # CrossEntropyLoss要求label格式为int,将Label格式转换为 int
    label = int(label)
    # 返回图像和对应标签
    return image, label

  def __len__(self):
    """
    步骤四:实现 __len__ 函数,返回数据集的样本总数
    """
    return len(self.data_list)

transform = Normalize(mean=[127.5], std=[127.5])
train_dataset = MyDataset('./mnist/train','./mnist/train/label.txt', transform)
test_dataset = MyDataset('./mnist/val','./mnist/val/label.txt', transform)
print('train_dataset images: ',len(train_dataset), 'test_dataset images: ',len(test_dataset))
train_dataset images:  60000 test_dataset images:  10000
以上代码定义了一个继承自 datasets.Dataset 的数据集类 MyDataset ,并根据数据集特性定制了 __init____getitem____len__ 三个函数。
  • __init__ 函数中,读取了图像数据路径和对应的标签,并一一对应存放在列表 data_list 中。

  • __getitem__ 函数中定义了图像按照指定索引读取图像,并对图像数据和标签进行预处理的过程,最终返回处理后图像 image 和对应的标签 label

  • __len__ 函数中返回该数据集包含的数据量,在这个案例中即列表 data_list 的长度。

数据预处理#

一般来说,数据在输入模型前,都会经过预处理阶段,将数据转换成模型所需要的格式与分布;同时为了提升训练模型的泛化性,也会对数据样本做一定的数据增强操作。
本节以图像数据为例,介绍数据预处理的方法。

内置数据处理方法#

datasets.transforms 中内置了一些数据预处理方法,可以通过以下代码查看:
from neurai import datasets

print('图像数据处理方法:', datasets.transforms.__all__)
图像数据处理方法: ['Compose', 'ToArray', 'Normalize', 'CenterCrop', 'RandomCrop', 'Resize', 'Pad', 'RandomHorizontalFlip', 'RandomVerticalFlip', 'PoissonEncoder', 'DirectEncoder', 'LatencyEncoder']
包括图像Resize,归一化,中心裁剪,随机裁剪等操作。
这些数据预处理方法既可以单独调用,也可以进行组合,多个函数一起调用。
  • 单个使用

from neurai.datasets.transforms import Resize
# 初始化一个调整图像大小的预处理方法
transform = Resize(size=28, interpolation='bicubic')
  • 多个组合使用

使用 Compose 组合多种预处理方法
from neurai.datasets.transforms import Compose, RandomHorizontalFlip, Resize
# 定义一个预处理组合方法,包含水平翻转和调整图像大小两个功能
transform = Compose([RandomHorizontalFlip(), Resize(size=28, interpolation='bicubic')])

数据集中应用数据预处理操作——内置函数#

  • 在内置数据集上应用

在初始化内置数据集时,将定义好的数据预处理方法通过 transform 字段传入即可
train_dataset = datasets.MNIST(root = './', train=True, transform=transform)
  • 在自定义数据集上应用

对于自定义的数据集,可以在数据集中将定义好的数据处理方法传入 __init__ 函数,将其定义为自定义数据集类的一个属性,然后在 __getitem__ 中将其应用到图像上,如下述代码所示:
import os
import cv2
from neurai import datasets
from neurai.datasets.transforms import Normalize

class MyDataset(datasets.Dataset):
  """
  步骤一:继承 datasets.Dataset 类
  """
  def __init__(self, root, label_path, transform=None):
    """
    步骤二:实现 __init__ 函数,初始化数据集,将样本和标签映射到列表中
    """
    super().__init__()
    self.data_list = []
    with open(label_path,encoding='utf-8') as f:
      for line in f.readlines():
          image_path, label = line.strip().split('\t')
          image_path = os.path.join(root, image_path)
          self.data_list.append([image_path, label])
    # 传入定义好的数据处理方法,作为自定义数据集类的一个属性
    self.transform = transform

  def __getitem__(self, index):
    """
    步骤三:实现 __getitem__ 函数,定义指定 index 时如何获取数据,并返回单条数据(样本数据、对应的标签)
    """
    image_path, label = self.data_list[index]
    image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    image = image.astype('float32')[...,None]
    # 应用数据处理方法到图像上
    if self.transform is not None:
        image = self.transform(image)
    label = int(label)
    return image, label

  def __len__(self):
    """
    步骤四:实现 __len__ 函数,返回数据集的样本总数
    """
    return len(self.data_list)

transform = Compose([RandomHorizontalFlip(),
                     Resize(size=28, interpolation='bicubic'),
                     Normalize(mean=[127.5], std=[127.5])])
train_dataset = MyDataset('./mnist/train','./mnist/train/label.txt', transform)

数据集中应用数据预处理操作——非内置函数#

除了内置的预处理函数,也可以根据实际需求在自定义数据集中设置数据预处理方法
以下代码使用opencv库完成前面相同的功能
import os
import cv2
from neurai import datasets
from neurai.datasets.transforms import Normalize

class MyDataset(datasets.Dataset):
  """
  步骤一:继承 datasets.Dataset 类
  """
  def __init__(self, root, label_path):
    """
    步骤二:实现 __init__ 函数,初始化数据集,将样本和标签映射到列表中
    """
    super().__init__()
    self.data_list = []
    with open(label_path,encoding='utf-8') as f:
      for line in f.readlines():
          image_path, label = line.strip().split('\t')
          image_path = os.path.join(root, image_path)
          self.data_list.append([image_path, label])

  def __getitem__(self, index):
    """
    步骤三:实现 __getitem__ 函数,定义指定 index 时如何获取数据,并返回单条数据(样本数据、对应的标签)
    """
    image_path, label = self.data_list[index]
    image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    # 应用数据处理方法到图像上
    image = cv2.resize(image, (14,14), interpolation=cv2.INTER_CUBIC)
    image = cv2.flip(image,flipCode=1)
    image = (image.astype('float32')-127.5)/127.5

    image = image[...,None]
    label = int(label)

    return image, label

  def __len__(self):
    """
    步骤四:实现 __len__ 函数,返回数据集的样本总数
    """
    return len(self.data_list)

train_dataset = MyDataset('./mnist/train','./mnist/train/label.txt')
在以上自定义数据集中,在 __getitem__ 函数中对数据直接使用opencv库进行预处理,本方案相对前一种方案灵活性更低,但更便于对需要特殊处理的数据集进行预处理。

迭代读取数据集#

直接迭代读取 Dataset 的方式虽然可实现对数据集的访问,但这种方案在训练或批量测试中需要手动将读取到的数据进行分批次(batch)操作。在本框架中,可以使用 datasets.DataLoader 对数据集进行读取,并且可以自动完成划分batch的工作。
train_loader = datasets.DataLoader(train_dataset, batch_size=64, shuffle=True, drop_last=True)
for batch_id, data in enumerate(train_loader):
  images, labels = data
  print("batch_id: {}, 训练数据shape: {}, 标签数据shape: {}".format(batch_id, images.shape, labels.shape))
  break
batch_id: 0, 训练数据shape: (64, 28, 28, 1), 标签数据shape: (64,)
以上代码初始化了一个数据读取器 train_loader ,用于加载训练数据集 train_dataset 。这个例子中使用的几个字段含义如下:
  • batch_size :每个batch包含的样本数量,在这里 batch_size=64 表示一个batch中包含64个数据样本。

  • shuffle :样本乱序,在这里 shuffle=True 表示在取数据时打乱数据样本顺序。

  • drop_last :丢弃不完整的批次样本,示例中 drop_last=True 表示丢弃因数据集样本数不能被 batch_size 整除而产生的最后一个不完整的 batch 样本。

定义好数据读取器之后,便可用 for 迭代读取批次数据,用于模型训练了。