petastorm
uber/petastorm加载项目详情…
本应用为开源项目，仅供学习研究，请遵守其开源协议。
微信扫码咨询客服
POMOAI — 让开源 AI 即刻运行
加载项目详情…
本应用为开源项目，仅供学习研究，请遵守其开源协议。
图1：Uber ATG（Advanced Technologies Group）开发的 Petastorm，是大规模深度学习数据访问的事实标准之一
2018 年，Uber 自动驾驶团队（ATG）的工程师们面临一个尴尬的局面：他们的感知模型每天需要处理 PB 级别的训练数据，这些数据存储在 Apache Parquet 格式中——一种在数据仓库领域极为高效的分析型存储格式。然而，TensorFlow 和 PyTorch 这样的深度学习框架根本不认 Parquet，工程师们只能先把 Parquet 转成 TFRecord 或 LMDB 等框架原生格式，这个 ETL 过程往往需要数小时，还容易出错。更糟糕的是，每次数据集更新都要重新跑一遍 ETL，数据新鲜度严重滞后。
Petastorm 就是为解决这个痛点而生的。它的核心思路非常优雅：让深度学习框架直接读取 Parquet，不需要任何格式转换。这一看似简单的改变，带来了数据管道从「ETL 驱动」到「数据驱动」的范式转变。
Petastorm 是 Uber ATG（Advanced Technologies Group）开源的数据访问库，专门解决大规模训练数据的高效读取问题。它支持将 Apache Parquet 格式的数据集直接供 TensorFlow、PyTorch 和 PySpark 使用，同时保持纯 Python 接口的易用性。
从技术定位上看，Petastorm 处于数据工程层和模型训练层的交界处。它不是数据处理工具（不负责生成 Parquet），而是一个高效的数据供给管道——像一个精心设计的「数据超市」，让训练进程可以高速、随机、多线程地「选购」数据。
Petastorm 的数据流设计非常精妙，可以概括为以下几个关键环节：
Apache Parquet 按 Row Group（行组）组织数据，每个 Row Group 默认包含 1000 行左右。Petastorm 的 make_reader 函数以 Row Group 为粒度进行并行读取，支持多线程、多进程甚至分布式并发消费。单个节点的并发度由 workers_count 参数控制，默认 10 个 worker 线程。
Petastorm 定义了 Unischema 作为跨框架的数据Schema抽象层。传统数据集中，字段类型在不同框架里有不同表达（如 TensorFlow 的 tf.float32 vs PyTorch 的 torch.float32）。Unischema 将所有字段类型统一管理，自动适配到目标框架的原生类型，开发者只需要定义一次 Schema，TensorFlow 和 PyTorch 都能直接使用，无需任何手写类型映射。
from petastorm.unischema import Unischema, UnischemaField
from petastorm.codecs import ScalarCodec, NdarrayCodec, CompressedImageCodec
from pyspark.sql.types import IntegerType, StringType
ImageSchema = Unischema('ImageSchema', [
UnischemaField('image_id', np.int32, (), ScalarCodec(IntegerType()), False),
UnischemaField('image', np.uint8, (None, None, 3), CompressedImageCodec('png'), False),
UnischemaField('label', np.int32, (), ScalarCodec(StringType()), False),
])
Petastorm 内置了两级缓存机制：本地磁盘缓存（local-disk）和空操作缓存（null）。在分布式训练场景中，同一数据 Row Group 可能被多个训练进程反复读取，开启本地缓存后，同一机器上所有进程共享磁盘缓存，显著减少对 HDFS/S3 的重复 I/O 请求。缓存大小由 cache_size_limit 控制，默认不设上限。
Petastorm 支持可插拔的数据编解码器。内置的 CompressedImageCodec 可以直接在 Parquet 中存储 JPEG/PNG 压缩图片，读取时自动解码为 numpy 数组，大幅节省存储空间。此外还有 NdarrayCodec 支持任意维度数组，满足点云、雷达、音频等多模态数据的存储需求。
Petastorm 对 PyTorch 的支持尤为成熟。通过 petastorm.pytorch.DataLoader，可以将 Petastorm Reader 无缝封装为 PyTorch 原生的 DataLoader，这意味着所有 PyTorch 生态工具（分布式训练、混合精度、Apex 等）都可以直接使用 Petastorm 管理的数据。
from petastorm.pytorch import DataLoader
from petastorm import make_reader
with make_reader('hdfs://path/to/parquet/') as reader:
with DataLoader(reader, batch_size=256) as loader:
for batch in loader:
# batch 已经是 PyTorch tensors，直接送入模型
outputs = model(batch)
值得注意的是，Petastorm 内置了 PyTorch 数据类型的自动提升逻辑——将 PyTorch 不原生支持的 int8、uint16、uint32 等类型自动转换为 int32/int64/uint8，避免烦人的类型不匹配错误。
Petastorm 原生支持数据分片（Sharding），这是分布式多 GPU 训练的关键功能。通过 cur_shard 和 shard_count 参数，可以将整个数据集划分为 N 个互不重叠的子集，每个 GPU 训练进程只消费属于自己的那一份数据，确保数据不重复、不遗漏。结合 PyTorch 的 DistributedDataParallel 或 TensorFlow 的 MultiWorkerMirroredStrategy，可以轻松构建分布式训练管道。
此外，Petastorm 支持 NGram 采样——从时间序列数据中提取连续的 N 个时间步，这在自动驾驶轨迹预测、强化学习等场景中非常有用。
Petastorm 本身是一个纯 Python 库，通过 pip install petastorm 即可完成安装。其生产环境使用通常需要以下基础设施：
Docker 方式提供了开箱即用的环境，但该镜像是为开发调试设计的（包含多个 Python 版本和完整依赖），镜像体积较大，不适合直接用于生产训练。
上手路径建议：从 examples/hello_world 开始，先用纯 Python 脚本跑通端到端流程，再根据需要引入 PySpark 数据生成和分布式训练。
Petastorm 并不是万能的，需要注意以下几点：
pyarrow>=6.0.1），使用过新或过旧的版本组合可能导致兼容性问题Petastorm 解决了 ML 基础设施中一个被长期忽视的「最后一公里」问题——数据格式转换。2018 年该项目发布时，业界普遍的做法是先 ETL 转格式再训练，而 Petastorm 证明了「直接读」不仅可行，而且效率更高。
其影响力体现在几个方面：
截至目前，该项目在 GitHub 拥有 1889 颗星，被 fork 284 次，181 个 open issues，反映了其活跃的社区关注度。
Petastorm 的核心价值在于消除了 ML 训练中的格式转换壁垒，让数据工程师和算法工程师可以用同一套 Parquet 数据高效协作。它不是炫技的框架，而是一款扎实的工程基础设施——解决了真实的大规模训练痛点，用代码质量说话。如果你正在构建大规模深度学习训练管线，且数据存储在 Parquet 格式，Petastorm 几乎是不二之选。