将S3中Parquet格式训练数据输入TensorFlow/Keras模型的最优方法是什么

优先选择Uber开源的Petastorm方案，其次是TensorFlow IO方案，完全不建议用Pandas分批读取的方案。

各方案适配大规模训练的优劣对比 #

• 方案1：Pandas分批读取
  仅适合超小数据集做逻辑验证，完全无法适配大规模训练需求。核心问题在于Pandas为单进程读取，所有批次数据都需要加载到Driver节点内存，数据量稍大就会出现OOM或者读取速度跟不上模型训练吞吐的问题，也完全不支持后续扩展分布式训练，直接排除即可。
• 方案2：TensorFlow IO APIs
  性能优于Pandas读取方案，但在Databricks场景下适配性较差。你提到的需要拷贝数据到本地的问题其实可以避免，TensorFlow IO本身支持直接通过S3协议读取远端对象存储的Parquet文件，不需要本地转存。但它对Parquet的复杂列类型兼容差，也无法利用Databricks集群的分布式算力做数据预处理，数据规模到TB级以上时读取瓶颈会非常明显。
• 方案3：Petastorm方案
  是三个方案中最适配当前场景的选择，你担心的需要拷贝数据到本地的问题也不存在：Petastorm原生支持直接读取S3上的Parquet数据集，还能完美对接Databricks的Spark运行环境，可以直接将Spark读取的Parquet DataFrame作为数据源喂给Keras模型，不需要额外做数据转存。
  针对大规模训练的核心优势非常明显：
  • 支持分布式数据加载，自动调度集群Worker节点完成数据分片读取、预处理操作，不会将读取压力全部集中在Driver节点
  • 原生适配TensorFlow/Keras的Dataset接口，支持数据加载和模型训练并行执行，不会出现数据供不上训练的问题
  • 后续如果要扩展分布式训练（比如对接Databricks上的分布式TensorFlow、Horovod组件），Petastorm生成的数据集可以直接无缝接入，几乎不需要修改核心代码

落地建议 #

在Databricks环境下使用时，直接通过spark.read.parquet("s3://<你的数据存储路径>")读入S3上的Parquet数据生成Spark DataFrame，再调用Petastorm的make_spark_dataset_converter接口即可直接转换为Keras可用的数据集，代码实现比另外两个方案更简洁。

内容的提问来源于stack exchange，提问作者exAres