FFRecord | 高性能训练数据格式

High-Flyer    July 27, 2020

在之前的文章《幻方萤火 | 高速读写文件系统 3FS》中提到,幻方AI自研了一套适合深度学习模型训练场景的文件读写系统,3FS,能提供高性能的批次数据读取,提高模型的训练效率。对于用户而言,使用3FS非常简单,只需要掌握我们封装设计的高性能数据格式,FFRecord,将数据存入萤火集群中即可。

那么FFRecord和一般的数据格式有什么不同?它应该要有哪些设计要求呢?本期文章将为大家分享FFRecord设计背后的故事,讲述幻方AI设计FFrecord的理念,展现FFRecord的高性能和便捷。

项目地址https://github.com/HFAiLab/ffrecord

概述

我们知道,3FS专门针对模型批量读取样本数据这个场景进行了深度的优化,和一般的文件系统不同,3FS 文件系统有如下的一些特点:

  1. 支持高吞吐的随机批量读取 (batch read)
  2. 大量打开、关闭小文件的开销比较大

如果要充分利用 3FS 文件系统的高效读取性能,我们希望读取的样本数据格式应该满足以下两个条件:

  1. 能够知道每一条样本的文件偏移量,方便做随机批量读取
  2. 避免大量的小文件,所有的样本数据在一个或多个大文件里

然而,目前网上的开源文件格式,比如 TFRecord 等并不完全满足上面的两个条件。TFRecord 能够避免大量的小文件,但是只支持顺序读,而且每次只能读取一条样本。它在不需要随机访问样本的情况下能够获得优秀的性能;但在很多机器学习的场景下,用户训练模型时每一步需要随机采样多个样本。

为此,我们专门设计了 FFRecord (Fire Flyer Record) 文件格式,帮助用户更加方便的享受 3FS 高速批量读取带来的效率提升。

以下是 FFRecord 文件格式的存储 layout:

+-----------------------------------+---------------------------------------+
|         checksum                  |             N                         |
+-----------------------------------+---------------------------------------+
|         checksums                 |           offsets                     |
+---------------------+---------------------+--------+----------------------+
|      sample 1       |      sample 2       | ....   |      sample N        |
+---------------------+---------------------+--------+----------------------+

在 FFRecord 文件格式中,我们会把每一条样本的数据做序列化按顺序写入,同时我们会存储每一条样本在文件中的偏移量和对应的 crc32 校验和,方便我们做随机读取和校验数据。相比于 TFRecord 等格式,我们没有直接给用户提供序列化和反序列化的接口,而是把这一步交给用户,用户可以用自己喜欢的方式 (比如 pickle)去做序列化或者反序列化。

I/O 接口

接下来我们介绍一下 FFRecord 的使用方式。FFRecord 文件的读写分别可以通过 FileReaderFileWriter 对象来完成。以下是一个简单的使用示例:

数据写入

from ffrecord import FileWriter

samples = [i for i in range(100)]  # 准备写入的样本
fname = 'test.ffr'
n = len(samples)                   # 样本的数量
writer = FileWriter(fname, n)

for i in range(n):
    data = pickle.dumps(samples[i])  # 序列化
    writer.write_one(data)

writer.close()

FileWriter.write_one 方法会把传入的数据写到文件的末尾,然后在文件的头部写入这条数据在文件中的偏移量和对应的 crc32 校验和

数据读取

from ffrecord import FileReader

fname = 'test.ffr'
reader = FileReader(fname, check_data=True)
print(f'Number of samples: {reader.n}')

indices = [3, 6, 0, 10]      # 样本的索引
data = reader.read(indices)  # 返回样本的原始数据
samples = [pickle.loads(x) for x in data]  # 反序列化

FileReader.read 方法会先根据传入的索引找到每条样本在文件中的偏移量并计算出占用的字节数,然后调用 Linux Asynchronize I/O (AIO) 的接口同时处理多个样本的读取请求,把结果返回给用户。

Dataset 和 DataLoader

PyTorch 本身提供了 map style 的 Dataset 和 DataLoader 的接口,然而它们的接口对需要批量读取的场景不是很友好,PyTorch Dataset 需要实现 __getitem__ 方法,该方法传入一个样本的索引并返回一条样本,这样如果我们希望批量读取多个样本就会比较困难。为了让 PyTorch 用户更加方便的使用 FFRecord 读取训练样本,我们提供了定制的 FFDataset 和 FFDataLoader。FFDataset 同样需要实现一个 __getitem__ 方法,但这个方法接受一个 batch 的样本索引作为输入,返回一整个 batch 的样本。FFDataset 可以搭配 FFDataLoader 一起使用,实现多进程异步加载数据。

以下是一个简单的使用示例:

from ffrecord.torch import Dataset as FFDataset
from ffrecord.torch import DataLoader as FFDataLoader

class ImageNet(FFDataset):
    def __init__(self, fname, transform=None):
        self.reader = FileReader(fname, check_data=True)
        self.transform = transform

    def __len__(self):
        return self.reader.n

    def __getitem__(self, indices, data):
        data = self.reader.read(indices)
        samples = []
        for bytes_ in data:
            img, label = pickle.loads(bytes_)
            if self.transform:
                img = self.transform(img)
            samples.append((img, label))
        return samples

dataset = ImageNet(fname, transform)
loader = FFDataLoader(dataset, batch_size=32, shuffle=True, num_workers=8)

性能测试

我们在 3FS 存储集群上进行了性能测试,以读取 ImageNet 数据集为例,比较 大量小文件 + PyTorch DataLoader 和 FFRecord + FFDataLoader的性能。为了模拟在分布式训练模型时候的读取模式,我们会起 8 个互相独立的进程,每个进程有自己的 DataLoader,然后每个 DataLoader 会再开 16 个子进程 (num_workers = 16),batch size 设为128,然后我们统计每个 DataLoader 读取 1000 个 batch 的时间并取平均。

image.png

上图展现了我们实验得到的结果,从图中我们可以看到使用了 FFRecord + FFDataLoder 后速度能够得到明显的提升。


本文作者: High-Flyer


您可以转载、不违背作品原意地摘录及引用本技术博客的内容,但必须遵守以下条款: 署名 — 您应当署名原作者,但不得以任何方式暗示幻方为您背书,亦不会对幻方的权利造成任何负面影响。 非商业性使用 — 您不得将本技术博客内容用于商业目的。 禁止演绎 — 如果基于该内容改编、转换、或者再创作,您不得公开或分发被修改内容,该内容仅可供个人使用。