Petastorm

.. image:: https://github.com/uber/petastorm/actions/workflows/unittest.yml/badge.svg?branch=master :target: https://github.com/uber/petastorm/actions/workflows/unittest.yml :alt: 构建状态

.. image:: https://codecov.io/gh/uber/petastorm/branch/master/graph/badge.svg :target: https://codecov.io/gh/uber/petastorm/branch/master :alt: 代码覆盖率

.. image:: https://img.shields.io/badge/License-Apache%202.0-blue.svg :target: https://img.shields.io/badge/License-Apache%202.0-blue.svg :alt: 许可证

.. image:: https://badge.fury.io/py/petastorm.svg :target: https://pypi.org/project/petastorm :alt: 最新版本

.. inclusion-marker-start-do-not-remove

.. contents::

Petastorm 是 Uber ATG 开发的一款开源数据访问库。该库支持直接从 Apache Parquet 格式的数据集中进行单机或分布式深度学习模型训练与评估。Petastorm 支持流行的基于 Python 的机器学习框架，如 Tensorflow <http://www.tensorflow.org/>、PyTorch <https://pytorch.org/> 和 PySpark <http://spark.apache.org/docs/latest/api/python/pyspark.html>_，同时也可在纯 Python 代码中使用。

文档网站： <https://petastorm.readthedocs.io>_

安装

.. code-block:: bash

pip install petastorm

petastorm 包定义了几项不会自动安装的额外依赖。这些附加组件包括：tf、tf_gpu、torch、opencv、docs 和 test。

例如，要安装 GPU 版本的 TensorFlow 和 OpenCV，可以使用以下 pip 命令：

.. code-block:: bash

pip install petastorm[opencv,tf_gpu]

生成数据集

使用 Petastorm 创建的数据集以 Apache Parquet <https://parquet.apache.org/>_ 格式存储。除了 Parquet 的 schema 外，Petastorm 还存储更高层次的 schema 信息，将多维数组原生地集成到 Petastorm 数据集中。

Petastorm 支持可扩展的数据编解码器。这些编解码器允许用户使用标准的数据压缩格式（如 jpeg、png），也可以实现自定义的编解码器。

数据集的生成通过 PySpark 完成。PySpark 原生支持 Parquet 格式，因此可以在单机或 Spark 计算集群上轻松运行。以下是一个将一些随机数据写入表中的极简示例。

.. code-block:: python

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 row_generator(x): """返回生成数据集中的单条记录。这里以返回一组随机值为例。""" return {'id': x, 'image1': np.random.randint(0, 255, dtype=np.uint8, size=(128, 256, 3)), 'array_4d': np.random.randint(0, 255, dtype=np.uint8, size=(4, 128, 30, 3))}

def generate_petastorm_dataset(output_url='file:///tmp/hello_world_dataset'): rowgroup_size_mb = 256

   spark = SparkSession.builder.config('spark.driver.memory', '2g').master('local[2]').getOrCreate()
   sc = spark.sparkContext

   # 包装数据集物化部分。它会负责设置 Spark 环境变量，并保存 Petastorm 特有的元数据。
   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()) \
           .coalesce(10) \
           .write \
           .mode('overwrite') \
           .parquet(output_url)

• HelloWorldSchema 是一个 Unischema 对象的实例。 Unischema 能够将其字段类型渲染为不同框架特定的格式，例如：Spark 的 StructType、TensorFlow 的 tf.DType 以及 NumPy 的 numpy.dtype。
• 要定义数据集字段，需要为每个 Unischema 字段指定类型、形状、编解码器实例，以及该字段是否可为空。
• 我们使用 PySpark 来写入输出的 Parquet 文件。在这个例子中，我们是在本地机器上启动 PySpark（.master('local[2]')）。当然，对于更大规模的数据集生成，我们需要一个真正的计算集群。
• 我们用 materialize_dataset 上下文管理器来包裹 Spark 数据集生成代码。该上下文管理器负责在开始时配置行组大小，并在结束时写入 Petastorm 特有的元数据。
• 行生成代码应返回一个以字段名索引的 Python 字典。我们使用 row_generator 函数来实现这一点。
• dict_to_spark_row 将字典转换为 pyspark.Row 对象，同时确保符合 HelloWorldSchema 的 schema 要求（会检查形状、类型和是否可空）。
• 一旦有了 pyspark.DataFrame，我们就将其写入 Parquet 存储中。Parquet 的 schema 会自动从 HelloWorldSchema 中推导出来。

纯 Python API

petastorm.reader.Reader 类是用户代码访问来自 Tensorflow 或 Pytorch 等机器学习框架数据的主要入口。Reader 具有多种功能，例如：

• 选择性列读取
• 多种并行策略：线程、进程、单线程（用于调试）
• N-gram 读取支持
• 行过滤（行谓词）
• 打乱顺序
• 用于多 GPU 训练的分区
• 本地缓存

使用 petastorm.reader.Reader 类读取数据集非常简单，可以通过 petastorm.make_reader 工厂方法创建：

.. code-block:: python

from petastorm import make_reader

with make_reader('hdfs://myhadoop/some_dataset') as reader:
   for row in reader:
       print(row)

支持的 URL 协议包括 hdfs://... 和 file://...。

一旦 Reader 实例化后，就可以将其作为迭代器使用。

TensorFlow API

要将 Reader 连接到 TensorFlow 图中，可以使用 tf_tensors 函数：

.. code-block:: python

from petastorm.tf_utils import tf_tensors

使用 make_reader('file:///some/localpath/a_dataset') 作为读取器时： row_tensors = tf_tensors(reader) 使用 tf.Session() 作为会话时： 对于范围内的 3 次迭代： 打印(session.run(row_tensors))

或者，您也可以使用新的 tf.data.Dataset API；

.. code-block:: python

from petastorm.tf_utils import make_petastorm_dataset

使用 `make_reader('file:///some/localpath/a_dataset')` 作为读取器时：
    dataset = make_petastorm_dataset(reader)
    迭代器 = dataset.make_one_shot_iterator()
    张量 = 迭代器.get_next()
    使用 `tf.Session()` 作为会话时：
        样本 = sess.run(tensor)
        打印(sample.id)

PyTorch API

如 pytorch_example.py <https://github.com/uber/petastorm/blob/master/examples/mnist/pytorch_example.py>_ 中所示， 通过 PyTorch 读取 Petastorm 数据集可以通过适配器类 petastorm.pytorch.DataLoader 来完成， 该类允许提供自定义的 PyTorch 聚合函数和转换。

请确保已安装 torch 和 torchvision：

.. code-block:: bash

pip install torchvision

下面的极简示例假设已定义了 Net 类以及 train 和 test 函数，这些内容包含在 pytorch_example 中：

.. code-block:: python

import torch
from petastorm.pytorch import DataLoader

torch.manual_seed(1)
设备 = torch.device('cpu')
模型 = Net().to(设备)
优化器 = torch.optim.SGD(模型参数, lr=0.01, momentum=0.5)

def _transform_row(mnist_row):
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    返回 (transform(mnist_row['image']), mnist_row['digit'])


transform = TransformSpec(_transform_row, removed_fields=['idx'])

使用 DataLoader(make_reader('file:///localpath/mnist/train', num_epochs=10,
                            transform_spec=transform, seed=1, shuffle_rows=True), batch_size=64) 作为训练数据加载器时：
    train(model, 设备, 训练数据加载器, 10, 优化器, 1)
使用 DataLoader(make_reader('file:///localpath/mnist/test', num_epochs=10,
                            transform_spec=transform), batch_size=1000) 作为测试数据加载器时：
    test(model, 设备, 测试数据加载器)

如果您正在处理非常大的批量大小，并且不需要支持 Decimal/字符串类型，我们提供了 petastorm.pytorch.BatchedDataLoader，它可以使用 Torch 张量（CPU 或 CUDA）进行缓冲，从而显著提高吞吐量。

如果您的数据集大小可以容纳在系统内存中，您可以使用内存中的数据加载器 petastorm.pytorch.InMemBatchedDataLoader。此数据加载器只需读取一次数据集，并将数据缓存在内存中，以避免在多个 epoch 中进行额外的 I/O 操作。

Spark 数据集转换器 API

Spark 转换器 API 简化了从 Spark 到 TensorFlow 或 PyTorch 的数据转换过程。 输入的 Spark DataFrame 首先被物化为 Parquet 格式，然后作为 tf.data.Dataset 或 torch.utils.data.DataLoader 加载。

下面的极简示例假设已定义了一个编译好的 tf.keras 模型，以及一个包含特征列和标签列的 Spark DataFrame。

.. code-block:: python

from petastorm.spark import SparkDatasetConverter, make_spark_converter
import tensorflow.compat.v1 as tf  # pylint: disable=import-error

# 首先指定一个缓存目录。
# 该目录用于保存物化的 Spark 数据帧文件
spark.conf.set(SparkDatasetConverter.PARENT_CACHE_DIR_URL_CONF, 'hdfs:/...')

df = ... # `df` 是一个 Spark 数据帧

# 从 `df` 创建一个转换器
# 它会将 `df` 物化到缓存目录。
converter = make_spark_converter(df)

# 从 `converter` 创建一个 TensorFlow 数据集
使用 converter.make_tf_dataset() 作为数据集时：
    # `dataset` 是 `tf.data.Dataset` 对象
    # 如果需要，可以对数据集进行转换
    dataset = dataset.map(...)
    # 我们可以在该数据集上训练/评估模型
    model.fit(dataset)
    # 当退出上下文时，数据集的读取器会被关闭

# 删除数据帧的缓存文件。
converter.delete()

下面的极简示例假设已定义了 Net 类以及 train 和 test 函数，这些内容包含在 pytorch_example.py <https://github.com/uber/petastorm/blob/master/examples/mnist/pytorch_example.py>_ 中， 以及一个包含特征列和标签列的 Spark DataFrame。

.. code-block:: python

from petastorm.spark import SparkDatasetConverter, make_spark_converter

# 首先指定一个缓存目录。
# 该目录用于保存物化的 Spark 数据帧文件
spark.conf.set(SparkDatasetConverter.PARENT_CACHE_DIR_URL_CONF, 'hdfs:/...')

df_train, df_test = ... # `df_train` 和 `df_test` 是 Spark 数据帧
模型 = Net()

# 从 `df_train` 创建一个转换器_train
# 它会将 `df_train` 物化到缓存目录。（`df_test` 同理）
converter_train = make_spark_converter(df_train)
converter_test = make_spark_converter(df_test）

# 从 `converter_train` 创建一个 PyTorch 数据加载器
使用 converter_train.make_torch_dataloader() 作为训练数据加载器时：
    # `dataloader_train` 是 `torch.utils.data.DataLoader` 对象
    # 我们可以使用该数据加载器训练模型
    train(model, 数据加载器, ...)
    # 当退出上下文时，数据集的读取器会被关闭

# 对于 `converter_test` 也执行相同的操作
使用 converter_test.make_torch_dataloader() 作为测试数据加载器时：
    test(model, 数据加载器, ...)

# 删除数据帧的缓存文件。
converter_train.delete()
converter_test.delete()

使用 PySpark 和 SQL 分析 Petastorm 数据集

Petastorm 数据集可以使用 PySpark 读入 Spark DataFrame，在那里您可以使用各种 Spark 工具来分析和操作数据集。

.. code-block:: python

从 Parquet 文件创建一个数据帧对象

dataframe = spark.read.parquet(dataset_url)

显示模式

dataframe.printSchema()

统计总数

dataframe.count()

显示单列

dataframe.select('id').show()

SQL 可以用来查询 Petastorm 数据集：

.. code-block:: python

spark.sql( 'SELECT count(id) ' 'from parquet.file:///tmp/hello_world_dataset').collect()

完整的代码示例可在以下位置找到：pyspark_hello_world.py <https://github.com/uber/petastorm/blob/master/examples/hello_world/petastorm_dataset/pyspark_hello_world.py>_，

非 Petastorm Parquet 存储

Petastorm 也可以直接从 Apache Parquet 存储中读取数据。为此，请使用 make_batch_reader（而不是 make_reader）。下表总结了 make_batch_reader 和 make_reader 函数之间的区别。

================================================================== ===================================================== make_reader make_batch_reader ================================================================== ===================================================== 仅限 Petastorm 数据集（使用 materializes_dataset 创建） 任何 Parquet 存储（某些原生 Parquet 列类型 尚不支持。

---

读取器每次返回一条记录。 读取器返回一批记录。批次大小不固定，由 Parquet 行组 大小决定。

---

传递给 make_reader 的谓词按单行评估。 传递给 make_batch_reader 的谓词按批次评估。

---

可根据 filters 参数过滤 Parquet 文件。 可根据 filters 参数过滤 Parquet 文件。 ================================================================== =====================================================

故障排除

请参阅 故障排除_ 页面；如果您找不到答案，请提交工单_。

另请参阅

1. Gruener, R., Cheng, O., 和 Litvin, Y. (2018) 介绍 Petastorm：Uber ATG 的深度学习数据访问库。网址：https://eng.uber.com/petastorm/
2. QCon.ai 2019：“Petastorm：构建机器学习流水线的轻量级方法” https://www.infoq.com/presentations/petastorm-ml-pipelines/_。

.. _Troubleshooting: docs/troubleshoot.rst .. _ticket: https://github.com/uber/petastorm/issues/new .. _Development: docs/development.rst

如何贡献

我们更倾向于以 GitHub 拉取请求的形式接收贡献。请将拉取请求发送到 github.com/uber/petastorm 仓库。

• 如果您正在寻找可以贡献的内容，请查看 GitHub 问题 <https://github.com/uber/petastorm/issues>_ 并在相关问题下留言。
• 如果您有改进建议，或想报告一个错误但没有时间修复它，请 创建一个 GitHub 问题 <https://github.com/uber/petastorm/issues/new>_。

要贡献补丁：

• 尽可能将您的工作拆分为小型、单一功能的补丁。合并包含大量不相关功能的大规模更改会困难得多。
• 将补丁作为针对主分支的 GitHub 拉取请求提交。有关教程，请参阅 GitHub 关于分叉仓库和发送拉取请求的指南。
• 在拉取请求中包含对拟议更改的详细描述。
• 确保您的代码通过单元测试。您可以在此处找到运行单元测试的说明：这里 <https://github.com/uber/petastorm/blob/master/docs/development.rst>_。
• 为您的代码添加新的单元测试。

提前感谢您的贡献！

有关开发相关信息，请参阅 开发_。

.. inclusion-marker-end-do-not-remove Place contents above here if they should also appear in read-the-docs. Contents below are already part of the read-the-docs table of contents.

Petastorm 快速上手指南

Petastorm 是 Uber ATG 开发的开源数据访问库，支持直接从 Apache Parquet 格式的数据集中进行单机或分布式的深度学习模型训练与评估。它完美适配 TensorFlow、PyTorch 和 PySpark 等主流框架。

环境准备

• 操作系统: Linux, macOS, Windows
• Python 版本: 3.6+
• 核心依赖:
  • Apache Spark (PySpark): 用于生成数据集（可本地运行或集群运行）。
  • 深度学习框架：根据需求选择安装 TensorFlow 或 PyTorch。
• 存储支持: 支持本地文件系统 (file://) 和 HDFS (hdfs://)。

国内加速建议： 在安装 Python 依赖时，推荐使用清华或阿里镜像源以提升下载速度： pip install -i https://pypi.tuna.tsinghua.edu.cn/simple ...

安装步骤

1. 基础安装

仅安装核心功能：

pip install petastorm

2. 按需安装额外依赖

Petastorm 将特定框架的依赖作为“extras”提供，需手动指定。

• 安装 PyTorch 支持:

  pip install petastorm[torch]
  

• 安装 TensorFlow GPU 版及 OpenCV 支持:

  pip install petastorm[opencv,tf_gpu]
  

• 其他可选组件: tf (CPU 版 TF), docs, test 等。

基本使用

Petastorm 的工作流分为两步：生成数据集（使用 PySpark）和 读取数据（用于训练）。

第一步：生成 Parquet 数据集

使用 PySpark 将数据转换为 Petastorm 格式的 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

# 1. 定义数据架构 (Unischema)
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 row_generator(x):
    """生成单条数据记录"""
    return {'id': x,
            'image1': np.random.randint(0, 255, dtype=np.uint8, size=(128, 256, 3)),
            'array_4d': np.random.randint(0, 255, dtype=np.uint8, size=(4, 128, 30, 3))}

def generate_petastorm_dataset(output_url='file:///tmp/hello_world_dataset'):
    rowgroup_size_mb = 256

    # 初始化 Spark (本地模式)
    spark = SparkSession.builder.config('spark.driver.memory', '2g').master('local[2]').getOrCreate()
    sc = spark.sparkContext

    rows_count = 10
    
    # 使用 materialize_dataset 上下文管理器处理元数据
    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()) \
            .coalesce(10) \
            .write \
            .mode('overwrite') \
            .parquet(output_url)

# 执行生成
if __name__ == '__main__':
    generate_petastorm_dataset()

第二步：读取数据

方式 A：纯 Python / 通用迭代器

最基础的读取方式，适用于调试或自定义数据处理逻辑。

from petastorm import make_reader

# 支持 file:// 和 hdfs:// 协议
with make_reader('file:///tmp/hello_world_dataset') as reader:
    for row in reader:
        print(row.id, row.image1.shape)

方式 B：PyTorch 集成

直接使用 DataLoader 对接 PyTorch 训练循环，支持数据变换（Transform）和洗牌（Shuffle）。

import torch
from petastorm import make_reader
from petastorm.pytorch import DataLoader
from torchvision import transforms

# 定义数据变换逻辑
def _transform_row(row):
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5,), (0.5,))
    ])
    # 返回 (输入数据，标签)
    return (transform(row['image1']), row['id'])

from petastorm.transform import TransformSpec
transform_spec = TransformSpec(_transform_row, removed_fields=['array_4d'])

# 创建 DataLoader
with DataLoader(make_reader('file:///tmp/hello_world_dataset', 
                            num_epochs=1, 
                            shuffle_rows=True,
                            transform_spec=transform_spec), 
                batch_size=4) as train_loader:
    
    for data, target in train_loader:
        # 此处接入你的 PyTorch 训练代码
        print(f"Batch shape: {data.shape}")

方式 C：TensorFlow 集成

支持 tf.data.Dataset API，可无缝嵌入 TensorFlow 计算图。

from petastorm import make_reader
from petastorm.tf_utils import make_petastorm_dataset
import tensorflow as tf

with make_reader('file:///tmp/hello_world_dataset') as reader:
    # 转换为 tf.data.Dataset
    dataset = make_petastorm_dataset(reader)
    
    # 应用常见的 TF 数据预处理
    dataset = dataset.map(lambda x: (x.image1, x.id)).batch(4)
    
    iterator = iter(dataset)
    sample = next(iterator)
    print(sample)

方式 D：Spark DataFrame 直接转换

如果你已经有了 Spark DataFrame，可以使用转换器直接生成训练数据，无需手动编写 Parquet 生成代码。

from petastorm.spark import SparkDatasetConverter, make_spark_converter
import tensorflow.compat.v1 as tf

# 设置缓存目录 (必须配置)
spark.conf.set(SparkDatasetConverter.PARENT_CACHE_DIR_URL_CONF, 'file:///tmp/petastorm_cache')

# 假设 df 是一个已有的 Spark DataFrame
# converter = make_spark_converter(df) 

# 生成 TF 数据集
# with converter.make_tf_dataset() as dataset:
#     model.fit(dataset)
    
# 使用完毕后清理缓存
# converter.delete()