Petastorm简介
Petastorm是由Uber ATG(Advanced Technologies Group)开发的一个开源数据访问库。它的主要目的是使深度学习模型能够直接从Apache Parquet格式的大规模数据集进行训练和评估,无需复杂的数据预处理和转换步骤。Petastorm支持单机和分布式环境,兼容多种流行的Python机器学习框架,包括TensorFlow、PyTorch和PySpark等。
主要特性
Petastorm具有以下几个主要特性:
-
直接读取Parquet数据:支持从Parquet格式的数据集直接读取数据进行模型训练,无需额外的数据转换步骤。
-
兼容多种ML框架:可与TensorFlow、PyTorch、PySpark等主流机器学习框架无缝集成。
-
单机和分布式支持:既可用于单机环境,也支持在分布式集群上进行大规模训练。
-
高效的数据读取:提供多种并行读取策略,支持数据缓存,以提高数据加载效率。
-
灵活的数据处理:支持选择性列读取、N-gram读取、数据过滤、洗牌等功能。
-
易用的API:提供简洁直观的Python API,易于上手使用。
安装和使用
安装
Petastorm可以通过pip进行安装:
pip install petastorm
对于GPU支持或其他额外功能,可以指定相应的extras:
pip install petastorm[opencv,tf_gpu]
基本使用
使用Petastorm读取数据集的基本流程如下:
from petastorm import make_reader
with make_reader('hdfs://myhadoop/some_dataset') as reader:
for row in reader:
print(row)
这里的'hdfs://'和'file://'都是支持的URL协议。
与TensorFlow集成
Petastorm提供了与TensorFlow集成的API:
from petastorm.tf_utils import make_petastorm_dataset
with make_reader('file:///some/localpath/a_dataset') as reader:
dataset = make_petastorm_dataset(reader)
iterator = dataset.make_one_shot_iterator()
tensor = iterator.get_next()
with tf.Session() as sess:
sample = sess.run(tensor)
print(sample.id)
与PyTorch集成
同样,Petastorm也可以很方便地与PyTorch一起使用:
from petastorm.pytorch import DataLoader
with DataLoader(make_reader('file:///localpath/mnist/train', num_epochs=10,
transform_spec=transform, shuffle_rows=True),
batch_size=64) as train_loader:
train(model, device, train_loader, optimizer)
数据集生成
Petastorm使用PySpark来生成数据集,支持从各种数据源创建Parquet格式的数据集。以下是一个简单的示例:
import numpy as np
from pyspark.sql import SparkSession
from pyspark.sql.types import IntegerType
from petastorm.codecs import ScalarCodec, CompressedImageCodec, NdarrayCodec
from petastorm.etl.dataset_metadata import materialize_dataset
from petastorm.unischema import dict_to_spark_row, Unischema, UnischemaField
# 定义数据集schema
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),
])
# 生成数据集
def generate_petastorm_dataset(output_url='file:///tmp/hello_world_dataset'):
spark = SparkSession.builder.getOrCreate()
sc = spark.sparkContext
with materialize_dataset(spark, output_url, HelloWorldSchema, 256):
rows_rdd = sc.parallelize(range(10))\
.map(row_generator)\
.map(lambda x: dict_to_spark_row(HelloWorldSchema, x))
spark.createDataFrame(rows_rdd, HelloWorldSchema.as_spark_schema())\
.write\
.mode('overwrite')\
.parquet(output_url)
这个示例展示了如何定义一个自定义的数据集schema,并使用PySpark生成相应的Parquet数据集。
高级功能
除了基本的数据读取功能外,Petastorm还提供了许多高级特性:
-
选择性列读取:可以只读取需要的列,提高效率。
-
多种并行策略:支持线程、进程和单线程(用于调试)等多种并行读取策略。
-
N-gram读取:支持读取连续的多个样本。
-
行过滤:可以使用谓词函数对数据进行过滤。
-
数据洗牌:支持在读取时对数据进行随机洗牌。
-
多GPU训练分区:支持为多GPU训练自动分区数据。
-
本地缓存:可以将频繁访问的数据缓存到本地,提高读取速度。
Spark数据集转换API
Petastorm提供了Spark数据集转换API,简化了从Spark到TensorFlow或PyTorch的数据转换过程。以下是一个使用TensorFlow的示例:
from petastorm.spark import SparkDatasetConverter, make_spark_converter
import tensorflow as tf
# 配置缓存目录
spark.conf.set(SparkDatasetConverter.PARENT_CACHE_DIR_URL_CONF, 'hdfs:/...')
# 创建转换器
converter = make_spark_converter(df)
# 使用转换器创建TensorFlow数据集
with converter.make_tf_dataset() as dataset:
# 可以进行数据集转换
dataset = dataset.map(...)
# 训练模型
model.fit(dataset)
# 删除缓存的数据文件
converter.delete()
这种方法可以将Spark DataFrame直接转换为TensorFlow或PyTorch可用的数据格式,大大简化了数据准备流程。
性能优化
为了在大规模数据集上获得最佳性能,Petastorm提供了几种优化策略:
-
使用合适的行组大小:在生成数据集时,选择适当的行组大小可以显著影响读取性能。
-
利用缓存:对于反复访问的小数据集,使用内存缓存可以大幅提升性能。
-
选择性读取:只读取必要的列,减少I/O开销。
-
并行读取:根据硬件资源,调整并行度以最大化吞吐量。
-
数据预取:使用预取功能可以在处理当前批次数据时并行加载下一批次数据。
结语
Petastorm作为一个强大的数据访问库,为大规模深度学习任务提供了高效、灵活的数据加载解决方案。它bridged the gap between big data systems和深度学习框架,使得从海量数据直接训练复杂模型成为可能。无论是在学术研究还是工业应用中,Petastorm都是一个值得考虑的工具。
随着深度学习和大数据技术的不断发展,我们可以期待Petastorm在未来会有更多的改进和新特性,为AI领域的发展做出更大的贡献。对于那些需要处理大规模数据集的机器学习工程师和研究人员来说,深入学习和使用Petastorm无疑是一个明智的选择。
