转换图片分类数据集，以便与 Cloud TPU 搭配使用

---

本教程介绍如何使用图片分类数据转换器示例脚本将原始图片分类数据集转换为用于训练 Cloud TPU 模型的 TFRecord 格式。

与将每个图片作为单个文件读取相比，TFRecord 从 Google Cloud Storage 中读取大型文件会更加高效。您可以在使用 tf.data.Dataset 流水线的任何位置使用 TFRecord。

如需详细了解如何使用 TFRecord，请参阅以下 TensorFlow 文档：

• TFRecord 和 tf.train.Example
• tf.data.Dataset
• tf.data: Build TensorFlow 输入流水线
• PyTorch TFRecord 读取器和写入器

如果您使用 PyTorch 或 JAX 框架，并且没有将 Google Cloud 存储空间用于数据集存储，则可能无法从 TFRecord 受益。

转换概览

GitHub 上的数据转换器代码库中的图片分类文件夹包含 converter 脚本 image_classification_data.py 以及实现示例 simple_example.py，您可以复制并进行修改以执行自己的数据转换。

图片分类数据转换器示例定义了两个类：ImageClassificationConfig 和 ImageClassificationBuilder。这些类在 tpu/tools/data_converter/image_classification_data.py 中定义。

ImageClassificationConfig 是一个抽象基类。您可为 ImageClassificationConfig 设置子类，以定义实例化 ImageClassificationBuilder 所需的配置。

ImageClassificationBuilder 是用于图片分类数据集的 TensorFlow 数据集构建器。它是 tdfs.core.GeneratorBasedBuilder 的子类。 它会从数据集中检索数据示例，并将其转换为 TFRecord。系统会将 TFRecord 写入由 ImageClassificationBuilder 的 __init__ 方法的 data_dir 参数指定的路径。

在 simple_example.py 中，SimpleDatasetConfig 为 ImageClassificationConfig 设置子类，实现用于定义支持的模式、图片类数量的属性，以及生成包含图片数据和图片类的字典的示例生成器，用于数据集内的每个示例。

main() 函数创建随机生成的图片数据的数据集，并实例化一个 SimpleDatasetConfig 对象，指定类数和数据集在磁盘上的路径。接下来，main() 实例化一个 ImageClassificationBuilder 对象，并传入 SimpleDatasetConfig 实例。最后，main() 调用 download_and_prepare()。调用此方法时，ImageClassificationBuilder 实例使用由 SimpleDatasetConfig 实现的数据示例生成器来加载每个示例，并将其保存到一系列 TFRecord 文件中。

如需更详细的说明，请参阅分类转换器笔记本。

准备工作

修改数据转换示例以载入数据集

要将数据集转换为 TFRecord 格式，请为定义以下属性的 ImageClassificationConfig 类设置子类：

• num_labels - 返回图片类的数量
• supported_modes - 返回数据集支持的模式列表（例如：test、train 和 validate）
• text_label_map - 返回用于对文本类标签和整数类标签之间的映射进行建模的字典（SimpleDatasetConfig 不使用此属性，因为它不需要映射）
• download_path - 下载数据集的路径（SimpleDatasetConfig 不使用此属性，example_generator 从磁盘加载数据）

实现 example_generator 生成器函数。此方法必须生成包含每个示例的图片数据和图片类名称的字典。 ImageClassificationBuilder 使用 example_generator() 函数检索每个示例并将其写入 TFRecord 格式的磁盘。

运行数据转化示例

1. 使用以下命令创建 Cloud Storage 存储桶：

   注意：在以下命令中，将 bucket-name 替换为您要分配给存储桶的名称。

   gsutil mb -p ${PROJECT_ID} -c standard -l us-central1 gs://bucket-name
   

2. 使用 gcloud 命令启动 Compute Engine 虚拟机。

   $ gcloud compute tpus execution-groups create \
    --vm-only \
    --zone=us-central1-b \
    --name=imageclassificationconverter \
    --tf-version=2.5.0
   

   注意 ：如果您未连接到 Compute Engine 实例，则可以运行以下命令进行连接：

   gcloud compute ssh imageclassificationconverter --zone=us-central1-b 

   从现在开始，前缀 (vm)$ 表示您应该在 Compute Engine 虚拟机实例上运行该命令。

3. 安装必需的软件包。

   (vm)$ pip3 install opencv-python-headless pillow
   

4. 创建脚本使用的以下环境变量。

   (vm)$ export STORAGE_BUCKET=gs://bucket-name
   

   (vm)$ export CONVERTED_DIR=$HOME/tfrecords
   (vm)$ export GENERATED_DATA=$HOME/data
   (vm)$ export GCS_CONVERTED=$STORAGE_BUCKET/data_converter/image_classification/tfrecords
   (vm)$ export GCS_RAW=$STORAGE_BUCKET/image_classification/raw
   (vm)$ export PYTHONPATH="$PYTHONPATH:/usr/share/tpu/models"
   

5. 切换到 data_converter 目录。

   (vm)$ cd /usr/share/tpu/tools/data_converter
   

对伪造的数据集运行数据转换器

simple_example.py 脚本位于数据转换器示例的 image_classification 文件夹中。使用以下参数运行脚本会生成一组虚构图片，并将其转换为 TFRecord。

(vm)$ python3 image_classification/simple_example.py \
  --num_classes=1000 \
  --data_path=$GENERATED_DATA \
  --generate=True \
  --num_examples_per_class_low=10 \
  --num_examples_per_class_high=11 \
  --save_dir=$CONVERTED_DIR

在我们的某个原始数据集上运行数据转换器

1. 为原始数据的位置创建一个环境变量。

   (vm)$ export GCS_RAW=gs://cloud-tpu-test-datasets/data_converter/raw_image_classification
   

2. 运行 simple_example.py 脚本。

   (vm)$ python3 image_classification/simple_example.py \
   --num_classes=1000 \
   --data_path=$GCS_RAW \
   --generate=False \
   --save_dir=$CONVERTED_DIR
   

simple_example.py 脚本采用以下参数：

• num_classes 表示数据集中的类数。为了与 ImageNet 格式一致，我们在这里使用 1000。
• generate 确定是否要生成原始数据。
• data_path 是指应生成数据的路径（如果 generate=True），或存储原始数据的路径（如果 generate=False）。
• num_examples_per_class_low 和 num_examples_per_class_high 用于确定要为每个类生成多少个示例。脚本会为此范围内的示例生成一个随机数。
• save_dir 是指已保存 TFRecord 的保存位置。要在 Cloud TPU 上训练模型，数据必须存储在 Cloud Storage 上。 这可以位于 Cloud Storage 或虚拟机上。

重命名 TFRecord 并将其移动到 Cloud Storage

以下示例将转换后的数据用于 ResNet 模型。

1. 将 TFRecord 重命名为与 ImageNet TFRecord 相同的格式：

   (vm)$ cd $CONVERTED_DIR/image_classification_builder/Simple/0.1.0/
   (vm)$ sudo apt install rename
   

   (vm)$ rename -v 's/image_classification_builder-(\w+)\.tfrecord/$1/g' *

2. 将 TFRecord 复制到 Cloud Storage：

   (vm)$ gsutil -m cp train* $GCS_CONVERTED
   (vm)$ gsutil -m cp validation* $GCS_CONVERTED