文 / AI 工程师 Mike Bernico 和 Carlos Ezequiel,Google Cloud
在训练计算机视觉机器学习模型时,数据加载是一种常见的性能瓶颈,可能导致 GPU 或 TPU 资源在等待数据加载到模型时得不到充分利用。以高效 TensorFlow Record (TFRecord) 格式存储数据集是解决这些问题的好方法,只不过,创建 TFRecords 往往需要大量的复杂代码。
我们近期开源了 TensorFlow Recorder 项目(也就是 TFRecorder),如此一来,数据科学家、数据工程师或 AI/ML 工程师只需几行代码即可创建基于图像的 TFRecords。使用 TFRecords 对于创建高效 TensorFlow ML 流水线非常重要,而过去的创建方法非常繁琐。在 TFRecorder 之前,要大规模创建 TFRecords,必须编写一个数据流水线来解析结构化数据,从存储中加载图像并将结果序列化为 TFRecord 格式。TFRecorder 允许您直接从 Pandas dataframe 或 CSV 写入 TFRecords,无需编写任何复杂的代码。
以下是 TFRecoder 的示例,但我们先谈谈 TFRecords 的一些特定优势。
TFRecords 如何提供帮助
TFRecord 文件格式可以将数据存储在文件集中,每个文件都包含序列化为二进制记录的 Protocol Buffers 序列,可以非常高效地读取,有助于减少上述数据加载瓶颈。
在使用 TFRecord 格式的同时实现 预提取和并行交错,可以进一步提高数据加载性能。当模型在当前步骤上执行训练时,使用 prefetch 可以提前获取下一个训练步骤的数据,从而缩短每个模型训练步骤的时间,Parallel interleave 允许您读取多个 TFRecords 分片(TFRecord 文件的一部分),并对这些交错的数据流进行预处理。这能够减少读取训练批次所需的延迟,特别适用于从网络读取数据。
使用 TensorFlow Recorder
使用 TFRecorder 只需几行代码即可创建 TFRecord。工作原理如下。
import pandas as pd
import tfrecorder
df = pd.read_csv(...)
df.tensorflow.to_tfrecord(output_dir="gs://my/bucket")
TFRecorder 目前预期数据与 Google AutoML Vision 格式相同。
这种格式与 Pandas dataframe 或 CSV 格式类似:
split | image_uri | label |
---|---|---|
TRAIN | gs://my/bucket/image1.jpg | cat |
其中:
- split 可以取值 TRAIN、VALIDATION 和 TEST
- image_uri 指定图像文件的本地或 Google Cloud Storage 位置。
- label 可以是将被整数化的基于文本的标签,也可以是一个整数
将来,我们希望进一步扩展 TensorFlow Recorder,支持所有格式的数据。
这个示例虽然能够将上千个图像顺利转换成 TFRecords,但扩展到百万计的图像可能就难以处理。为了扩展到庞大的数据集,TensorFlow Recorder 提供了与 Google Cloud Dataflow 的连接,后者是一款无服务器 Apache Beam 流水线运行器。扩展到 DataFlow 仅需要多一点配置。
df.tensorflow.to_tfrecord(
output_dir="gs://my/bucket",
runner="DataFlowRunner",
project="my-project",
region="us-central1)
未来计划
我们希望您愿意尝试 TensorFlow Recorder。您可以从 GitHub 获取,或者直接 pip 安装 tfrecorder。
TensorFlow Recorder 刚刚面世,我们非常期待您的反馈、建议和 Pull Requests。