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

准备工作

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

您不必使用 TFRecord 为 TPU 的数据输入。使用 TFRecord 是 TensorFlow 配合 TPU 的首选方法,因为这种工具能够提高从 Google Cloud Storage 读取大文件的效率。您可以在使用 tf.data.Dataset 流水线的任何位置使用 TFRecord。

如需详细了解如何将 TFRecord 与 TensorFlow 数据集搭配使用,请参阅以下 TensorFlow 文档:

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

转换概览

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

图片分类数据转换器示例定义了两个类:ImageClassificationConfigImageClassificationBuilder。这些类在 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 中,SimpleDatasetConfigImageClassificationConfig 设置子类,实现用于定义支持的模式、图片类数量的属性,以及生成包含图片数据和图片类的字典的示例生成器,用于数据集内的每个示例。

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 存储分区:

    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
    
    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_lownum_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

清理

为避免因本教程中使用的资源导致您的 Google Cloud 帐号产生费用,请删除包含这些资源的项目,或者保留项目但删除各个资源。

  1. 断开与 Compute Engine 实例的连接(如果您尚未这样做):

    (vm)$ exit
    

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

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

    $ gcloud compute tpus execution-groups delete imageclassificationconverter \
      --zone=us-central1-b
    
  3. 通过运行 gcloud compute tpus execution-groups list 来验证虚拟机是否已删除。删除操作可能需要几分钟时间才能完成。如下所示的响应表明实例已成功删除。

    $ gcloud compute tpus execution-groups list --zone=us-central1-b
    
    Listed 0 items.
    
  4. 如下所示运行 gsutil,将 bucket-name 替换为您为本教程创建的 Cloud Storage 存储分区的名称:

    $ gsutil rm -r gs://bucket-name