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

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

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

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

使用 PyTorch 或 JAX 框架，并且未使用 Cloud Storage 数据集存储，您可能无法获得 TFRecords 同样的优势。

转换概览

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 文件中。

有关详细说明，请参阅 Classification Converter 笔记本。

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

要将数据集转换为 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 存储分区：

    gcloud storage buckets create gs://bucket-name --project=${PROJECT_ID} --location=us-central2

  2. 使用 gcloud 命令启动 Cloud TPU。

    $ gcloud compute tpus tpu-vm create tpu-name \
    --zone=us-central2-b \
    --accelerator-type=v4-8 \
    --version=tpu-vm-tf-2.17.0-pjrt

    命令标志说明

    zone
    拟在其中创建 Cloud TPU 的区域。
    accelerator-type
    加速器类型用于指定您要创建的 Cloud TPU 的版本和大小。 如需详细了解每个 TPU 版本支持的加速器类型，请参阅 TPU 版本。
    version
    Cloud TPU 软件版本。

  3. 使用 SSH 连接到 TPU：

    $ gcloud compute tpus tpu-vm ssh tpu-name --zone=us-central2-b

    连接到 TPU 后，shell 提示符将从 username@projectname 至 username@vm-name。

  4. 安装必需的软件包。

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

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

    (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"

  6. 切换到 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)$ gcloud storage cp train* $GCS_CONVERTED
(vm)$ gcloud storage cp validation* $GCS_CONVERTED

清理

  1. 断开与 Cloud TPU 的连接（如果尚未断开） 请按下列步骤操作：

    (vm)$ exit

    您的提示符现在应为 user@projectname，表明您位于 Cloud Shell 中。

  2. 在您的 Cloud Shell 中，运行 gcloud 以删除虚拟机资源。

    $ gcloud compute tpus tpu-vm delete tpu-name \
  --zone=us-central2-b

  3. 通过运行 gcloud compute tpus tpu-vm list 来验证虚拟机是否已删除。删除操作可能需要几分钟时间才能完成。如下所示的响应表明实例已成功删除。

    $ gcloud compute tpus tpu-vm list --zone=us-central2-b
     
    Listed 0 items.

  4. 如下所示，运行 gcloud CLI，将 bucket-name 替换为 您为本教程创建的 Cloud Storage 存储桶的名称：

    $ gcloud storage rm gs://bucket-name --recursive