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

本教程介绍如何使用图片分类数据转换器示例脚本将用于图片分类的原始数据集转换为 Cloud TPU Tensorflow 模型使用的 TFRecord 格式。GitHub 上的[图片分类代码库] (https://github.com/tensorflow/tpu/tree/master/tools/data_converter) 包含 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 -b on gs://bucket-name
    
  2. 使用 ctpu 命令创建虚拟机:

    ctpu up --vm-only \
      --zone=us-central1-b \
      --name=img-class-converter \
      --tf-version=2.2
    
    gcloud compute ssh img-class-converter --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

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

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

    (vm)$ cd $CONVERTED_DIR/image_classification_builder/Simple/0.1.0/
    (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

在生成的数据集上运行 ResNet

  1. 创建 TPU:

    (vm)$ export TPU_NAME=imageclassificationconverter
    (vm)$ ctpu up --zone=us-central1-b 
    --tpu-only
    --tpu-size=v2-8
    --name=${TPU_NAME}

  2. 为包含模型的存储分区目录创建一个环境变量。

    (vm)$ export MODEL_BUCKET=${STORAGE_BUCKET}/image_classification_converter
    
  3. 运行模型

    (vm)$ cd /usr/share/tpu/models/official/resnet
    (vm)$ python3 resnet_main.py --tpu=${TPU_NAME} 
    --data_dir=${GCS_CONVERTED}
    --config_file=configs/cloud/v3-8.yaml
    --model_dir=${MODEL_BUCKET}
    --num_train_images=10000
    --num_eval_images=10000