策划编辑 | Natalie
作者 | Uber ATG
翻译 | 无明
编辑 | Natalie
出处丨AI 前线
AI 前线导读: 近年来,深度学习在解决模式识别问题方面发挥了关键作用。Uber Advanced Technologies Group(ATG)使用深度学习来解决自动驾驶领域的各种问题,他们的很多模型需要来自传感器数十 TB 的训练数据。Uber ATG 的研究人员和工程师正在积极推动跨多个问题领域的自动驾驶技术,如感知、预测和规划。为了支持这些工作,他们致力于开发数据集存储解决方案,让研究人员更容易获得数据,从而可以专注于模型实验。本文将介绍 Petastorm,一个由 Uber ATG 开发的开源数据访问库。这个库可以直接基于数 TB Parquet 格式的数据集进行单机或分布式训练和深度学习模型评估。Petastorm 支持基于 Python 的机器学习框架,如 Tensorflow、Pytorch 和 PySpark,也可以直接用在 Python 代码中。
深度学习集群的搭建
即使是在现代硬件上训练深度模型也很耗时,而且在很多情况下,很有必要在多台机器上分配训练负载。典型的深度学习集群需要执行以下几个步骤:
- 一台或多台机器读取集中式或本地数据集。
- 每台机器计算损失函数的值,并根据模型参数计算梯度。在这一步通常会使用 GPU。
- 通过组合估计的梯度(通常由多台机器以分布式的方式计算得出)来更新模型系数。
图 1:在这个深度学习集群架构中,有三个计算节点使用中央数据集。
考虑到 GPU 的成本,很有必要提高 GPU 集群的利用率。经过调优的数据访问层可以确保用于训练的数据对 GPU 总是可用的,这样 GPU 就不会处于空闲状态。
简化模型架构研究
准备数 TB 来自多个数据源的同步数据通常很容易出错。我们希望为研究人员提供单个数据集,让他们可以处理各种任务,无需为每种任务创建新的数据集。
为此,需要遵守以下原则:
- 数据集需要包含研究人员可能用到的数据的超集,这样他们就可以为特定实验选择列和行的子集。
- 对数据集中的传感器数据的预处理应该保持在最低限度。我们鼓励研究人员进行实时的预处理,并
- 将其作为训练 / 评估程序的一部分。在很多情况下,这可以通过其他未被充分利用的 CPU 来完成。
在业界,深度学习应用程序的数据集存储通常分为两类:多文件和记录流式数据集。
多文件数据集
在这种情况下,每个张量 / 图像 / 标签集被保存在单独的文件(例如,PNG、JPEG、NPZ 和 CSV)中。整个数据集被存储为一个或多个文件系统目录,每个目录包含大量的文件。文件数量可能达到数百万个(例如,ImageNet 有 120 万个文件)。如果以这种格式存储,Uber ATG 的数据集将超过 1 亿个文件。
这种方法让用户可以随机访问数据集中任何行的任何列。但是,多次往返文件系统的成本很高,所以很难大规模实现,特别是在使用现代分布式文件系统时,如 HDFS 和 S3(这些系统通常针对大块数据的快速读取进行了优化)。
记录流式数据集
另外一种方式是将数据行的集合组合在一起,保存成一个或多个文件。例如,Tensorflow 使用 protobuf 文件(TFRecord)。其他流行的格式还包括 HDF5 和 Python pickle 文件。
这种方法适用于 HDFS 和 S3 文件系统。但是,查询特定列需要通过网络传输所有字段,然后丢弃未使用的数据。如果要查询单行,还需要自定义索引。
在评估了多个方案后,我们决定使用 Apache Parquet 存储格式,它在一定程度规避了上述两种方法的一些缺点:
- 便于进行大量连续读取(对 HDFS/S3 友好);
- 支持快速访问单个列;
- 在某些情况下允许更快的行查询;
- 与 Apache Spark 完美集成,可作为现成的查询 / 操作框架。
列式存储和 Apache Parquet
列式数据存储按照列(而不是行)的顺序来组织数据。例如,从自动驾驶车辆传感器记录的数据可能看起来像这样:
行和列存储之间的差异如下所示:
以列式顺序存储数据允许用户只加载列的子集,从而减少通过网络传输的数据量。对于收集来自自动驾驶车辆传感器的数据来说,这种好处是显而易见的:试想一下,如果你的实验只需要来自某个摄像头的图像,那么就可以从同一行的 10 张高分辨率图像中加载其中的一张。
Apache Parquet 是一种列式存储格式,近年来越来越流行。它得到了 Apache Spark、Apache Arrow 和其他开源项目的支持,并且非常适合用于进行简化模型架构研究。
Tensorflow 和 Pytorch 是深度学习社区常用的框架。这些框架本身并不支持 Parquet 存储访问,因此我们构建了 Petastorm 来填补这一空白。
Petastorm 简介
通常,一个数据集是通过连接多个数据源的记录而生成的。这个由 Apache Spark 的 Python 接口 PySpark 生成的数据集稍后将被用在机器学习训练中。Petastorm 提供了一个简单的功能,使用 Petastorm 特定的元数据对标准的 Parquet 进行了扩展,从而让它可以与 Petastorm 兼容。
有了 Petastorm,消费数据就像在 HDFS 或文件系统中创建和迭代读取对象一样简单。Petastorm 使用 PyArrow 来读取 Parquet 文件。
图 2:将多个数据源组合到单个表格结构中,从而生成数据集。可以多次使用相同的数据集进行模型训练和评估。
生成数据集
要使用 Petastorm 生成数据集,用户首先需要定义数据模式,也就是 Unischema。这是用户唯一需要定义模式的地方,Petastorm 会将它转换为其他框架所需的格式,例如 PySpark、Tensorflow 和 Python。
Unischema 的实例被序列化为 Parquet 存储元数据中的自定义字段,可以使用数据集的路径来读取它。
下面的示例演示了如何创建 Unischema 实例。必需的字段属性包括:字段名称、数据类型(使用 NumPy 数据类型来表示)、多维数组、用于数据编码 / 解码的 codec,以及一个表示字段是否可为空的布尔值。
HelloWorldSchema = Unischema('HelloWorldSchema', [
UnischemaField('id', np.int32, (), ScalarCodec(IntegerType()), False),
UnischemaField('image1', np.uint8, (128, 256, 3) CompressedImageCodec('png'), False),
UnischemaField('array_4d', np.uint8, (None, 128, 30, None), NdarrayCodec(), False),
])
我们使用 PySpark 来写入 Petastorm 数据集。下面的示例演示了如何使用 Petastorm 创建 1000 行数据。
rows_count = 10
with materialize_dataset(spark, output_url, HelloWorldSchema, rowgroup_size_mb):
rows_rdd = sc.parallelize(range(rows_count))\
.map(row_generator)\
.map(lambda x: dict_to_spark_row(HelloWorldSchema, x))
spark.createDataFrame(rows_rdd, HelloWorldSchema.as_spark_schema())\
.write \
.parquet('file:///tmp/hello_world_dataset')
- materialize_dataset 上下文管理器在开始时执行必要的配置,并在最后写入 Petastorm 元数据。输出的 URL 可以指向 HDFS 或文件系统的位置。
- rowgroup_size_mb 定义了 Parquet 行组的大小(以兆字节为单位)。
- row_generator 是一个返回与 HelloWorldSchema 匹配的 Python 字典的函数。
- dict_to_spark_row 根据 HelloWorldSchema 来验证数据类型,并将字典转换为 pyspark.Row 对象。
读取数据集
接下来,我们将概述如何使用 Python 代码以及在两个常用的机器学习框架(Tensorflow 和 Pytorch)中读取数据集。
Python
在 Python 代码中,可以直接使用 Reader 实例访问 Petastorm 数据集。Reader 实现了迭代器接口,所以读取数据很简单:
with Reader('file:///tmp/hello_world_dataset') as reader:
# Pure python
for sample in reader:
print(sample.id)
plt.imshow(sample.image1)
Tensorflow
下面的示例显示了如何将数据集流式传输到 Tensorflow。examples 是一个元组,它的键来自 Unischema,而值为 tf.tensor 对象:
with Reader('file:///tmp/hello_world_dataset') as reader:
tensor = tf_tensors(reader)
with tf.Session() as sess:
sample = sess.run(tensor)
print(sample.id)
plt.imshow(sample.image1)
在不久的将来,用户可以使用 tf.data.Dataset 接口来访问数据。
Pytorch
Petastorm 数据集可以通过适配器类 petastorm.pytorch.DataLoader 集成到 Pytorch 中,如下所示:
with DataLoader(Reader('file:///tmp/hello_world_dataset')) as train_loader:
sample = next(iter(train_loader))
print(sample['id'])
plt.plot(sample['image1'])
使用 Spark 分析数据集
Spark 本身支持 Parquet 数据格式,因此可以使用各种 Spark 工具来分析和操作数据集。下面的示例演示了如何将 Petastorm 数据集读取为 Spark RDD 对象:
rdd = dataset_as_rdd('file:///tmp/hello_world_dataset', spark,
[HelloWorldSchema.id, HelloWorldSchema.image1])
print(rdd.first().id)
标准的 PySpark 工具可用于处理 Petastorm 数据集。请注意,数据并不会被解码,而且只有在 Parquet 格式中具有相应原生表示的字段的值(例如标量)才有意义:
# Create a dataframe object from a parquet file
dataframe = spark.read.parquet(dataset_url)
# Show a schema
dataframe.printSchema()
# Count all
dataframe.count()
# Show a single column
dataframe.select('id').show()
可以使用 SQL 查询 Petastorm 数据集:
number_of_rows = spark.sql(
'SELECT count(id) '
'from parquet.`file:///tmp/hello_world_dataset`').collect()
Petastorm 的特性
Petastorm 提供了各种特性来支持自动驾驶算法的训练,包括行过滤、数据分片、shuffle、对字段子集的访问,以及对时间序列数据(n-gram)的支持。
典型数据集的结构包括:
- 在自动驾驶汽车测试运行期间收集的传感器数据的多个列,包括摄像头、激光定位器和雷达。
- 手动生成的标签,作为行的字段进行存储。
行数据按照时间顺序排序,并按照汽车的测试运行进行分组,行组大小通常在 30 到 100 范围内。
并行执行策略
Petastorm 提供了两种并行化数据加载和解码操作的策略:一种基于线程池,另一种基于进程池。策略的选择取决于所读取的数据类型。
通常,当行中包含编码的高分辨率图像时,应使用线程池策略。在这种情况下,大部分处理时间用于通过 C++ 代码来解码图像。这个时候不会持有 Python 全局解释器锁(GIL)。
当行很小时,使用进程池策略更合适。在这种情况下,大部分处理都通过 Python 代码来完成。这个时候必须并行运行多个进程,这样才能克服 GIL 导致的执行串行化。
n-gram
有些模型需要时间上下文,以便更好地解释环境或预测环境中参与者的未来行为。
如果底层的数据是按时间排列,Petastorm 就可以提供这样的时间上下文。如果向 Petastorm Reader 对象请求 n-gram,那么后续的行将被分组到单个训练样本中。
下图显示了长度为 3 的 n-gram 的分组。AV Log#0 和 AV Log#1 表示两种不同的车载记录:
图 3:在读取数据集时构造 n-gram。n-gram 不能跨 Parquet 行组。
请注意,n-gram 分组不能跨 Parquet 行组。在图 3 中,row-group 0 生成了三个 n-gram,而 row-group 1 只生成一个,另外三个来自 row-group 2。n-gram 节省了 IO 和 CPU 带宽,因为不需要进行磁盘数据复制,也不需要进行重复加载和解码。
n-gram 按照它们在数据集中出现的顺序来生成,因此用户需要让数据集中的顺序与访问模式保持一致。
shuffle
如果数据集支持 n-gram 访问模式,那么它的行数据时按时间戳排序的。Parquet 支持加载行组中的全部行。因此,数据将被加载到高度相关的样本组中(例如,从一辆自动驾驶车辆的摄像头获取的连续两张图像将非常相似)。连续样本之间的高相关性不是我们所期望的,它们会降低训练算法的性能。为了减少相关性,Petastorm 提供了 shuffle 的功能。
图 4:通过随机选择要加载的行组,然后将个体样本放入内存 shuffle 缓冲区来实现 shuffle。
Petastorm 从数据集中的随机选择一组行组。解码过的行被放入行 shuffle 缓冲区,然后从缓冲区中选择一个随机行返回给用户。
行谓词(过滤器)
对于在多个实验室中重用的数据集实例,能够有效地选择行子集是非常重要的。Petastorm 支持行谓词。Petastorm 行谓词利用了 Parquet 存储分区,只加载符合条件的列。
在未来,我们计划将 Parquet 的谓词下推(pushdown)功能集成到 Petastorm 中,以进一步加快查询。
行组索引
Petastorm 支持存储一个键与一组行组的映射。这种映射有助于快速查找符合特定条件的行组。在使用“行谓词”的地方,需要进行额外的过滤。
为分布式训练进行分片
在分布式训练环境中,每个进程通常负责训练数据的一个子集。一个进程的数据子集与其他进程的数据子集正交。Petastorm 支持将数据集的读时分片转换为正交的样本集。
图 5:Petastorm 将数据集的非重叠子集提供给参与分布式训练的不同机器。
本地缓存
Petastorm 支持在本地存储中缓存数据。当网络连接速度较慢或带宽很昂贵时,这会派上用场。
图 6:如果启用了本地缓存,每个会话仅下载一次数据。
在第一个时间段,从远程存储读取一组样本,并保存到本地缓存中。在随后的时间段,将从本地缓存中读取所有数据。
Petastorm 架构
Petastorm 的设计目标包括:
- 通过单数据模式定义进行数据的编码和解码。
- 为 ML 框架和纯 Python 代码提供可用的高数据加载带宽。
- 将 Apache Spark 作为分布式集群计算框架来生成数据集。
- 与 Python、ML 平台无关的 Petastorm 核心组件的实现。
- 呈现给 Tensorflow 和 PyTorch 框架的原生接口。
- etl 包实现了生成数据集的功能。
- Reader 是训练和计算代码使用的主要数据加载引擎。Reader 使用 Python 实现,不依赖任何 ML 框架(Tensorflow、Pytorch),并且可以通过 Python 代来实例化和使用。
- 为 Tensorflow 和 PyTorch 提供适配器。
- Unischema 可以被数据集生成和数据加载代码引用。
图 7:Petastorm 提供了支持数据集生成和读取的组件。Unischema 定义了可供两者使用的公共数据模式。
对 Parquet 的修复
当我们开始使用 Spark 写入 Parquet 数据集时遇到了一些麻烦。原因主要是数据的行大小,它包含了几个数兆字节的字段。我们的第一个问题是数据集中的行组比预期的要大得多,导致内存不足等问题。深入研究代码后发现,parquet-mr 在检查行组是否达到用户设置的目标大小之前,强制限制行组最少为 100 行。针对这个问题,Parquet 已经有一个相关的拉取请求,于是我们 fork 了这个代码库,并做了一些修改,让行组的大小符合我们的要求。
在解决了行组大小的问题后,我们发现,当我们尝试生成较小的行组或使用更大的字段时,Spark 作业会耗尽内存。通过深入研究生成的数据集,我们发现,在添加新字段或减少行组大小时,存储文件元数据的 Parquet 页脚大小显著增加了。
原来 Parquet 会为代表图像的巨大二进制字段或其他多维数组生成统计数据。由于 Parquet 在会页脚中保存这些字段的最小值和最大值,因此,如果行组大小足够小,那么页脚就会变得很大,直至无法全部放到内存中。这个问题在 parquet-mr 代码库中已经得到了解决,但是我们使用的是 Spark 2.1.0(依赖了 Parquet 1.8.1)。为了解决这个问题,我们升级了 Spark(Parquet 1.8.3 中已经修复了这个问题)。
下一步
下面我们重点介绍一下计划在不久的将来推出的一些改进:
减少 shuffle 的内存占用
大行组有助于提高 IO 利用率和数据加载速度。不过,它们也会增加连续样本之间的相关性。我们正在积极改进 shuffle 机制。
谓词下推支持
Pyarrow 将很快提供谓词下推支持。我们希望用它来实现更快的行过滤。
改进与 Spark 的集成
在 Spark 中访问 Petastorm 数据集时,某些操作似乎比预期花费更多的时间或内存。我们需要进一步调查 Parquet 库代码,以了解有效处理大型字段的其他细微差别。
额外的存储格式
Petastorm 抽象了底层存储格式。我们可以将 Parquet 以外的存储格式集成到 Petastorm 中,从而为实验和数据加载性能调整提供更大的自由。
GitHub 开源项目传送门:
Petastorm:
http://www.github.com/uber/petastorm
parquet-mr:
https://github.com/apache/parquet-mr
英文原文:
https://eng.uber.com/petastorm/