Petastorm: Uber开源的大规模深度学习数据访问库

petastorm

Petastorm简介

Petastorm是由Uber ATG(Advanced Technologies Group)开发的一个开源数据访问库。它的主要目的是使深度学习模型能够直接从Apache Parquet格式的大规模数据集进行训练和评估,无需复杂的数据预处理和转换步骤。Petastorm支持单机和分布式环境,兼容多种流行的Python机器学习框架,包括TensorFlow、PyTorch和PySpark等。

Petastorm架构

主要特性

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无疑是一个明智的选择。