下载、预处理和上传 COCO 数据集

COCO 是一种大规模的对象检测、细分和字幕数据集。使用 COCO 数据集的机器学习模型包括:

  • Mask-RCNN
  • Retinanet
  • ShapeMask

在 Cloud TPU 上训练模型之前,您必须先准备训练数据。由于系统在 TPU 设置完毕后开始对 Cloud TPU 计费,因此,最佳做法是设置 Compute Engine 虚拟机,准备数据集,然后设置 Cloud TPU。

本主题介绍如何为在 Cloud TPU 上运行的模型准备 COCO 数据集。只有在创建 Compute Engine 虚拟机后,才能准备 COCO 数据集。用于准备数据的脚本 download_and_preprocess_coco.sh 安装在虚拟机上,并且必须在虚拟机上运行。

通过运行 download_and_preprocess_coco.sh 脚本准备数据后,您可以启动 Cloud TPU 并运行训练。

准备 COCO 数据集

COCO 数据集将存储在您的 Cloud Storage 中,因此请设置一个存储分区变量来指定您创建的分区名称:

(vm)$ export STORAGE_BUCKET=gs://bucket-name
(vm)$ export DATA_DIR=${STORAGE_BUCKET}/coco

运行 download_and_preprocess_coco.sh 脚本,将 COCO 数据集转换为训练应用所需的一组 TFRecord (*.tfrecord)。

(vm)$ sudo bash /usr/share/tpu/tools/datasets/download_and_preprocess_coco.sh ./data/dir/coco

这会安装所需的库,然后运行预处理脚本。它会在您的本地数据目录中输出很多 *.tfrecord 文件。COCO 下载和转换脚本大约需要 1 小时才能完成。

将数据复制到 Cloud Storage 存储分区

将数据转换为 TFRecord 后,使用 gsutil 命令将其从本地存储空间复制到 Cloud Storage 存储分区。您还必须复制注释文件。这些文件有助于验证模型的性能。

(vm)$ gsutil -m cp ./data/dir/coco/*.tfrecord ${DATA_DIR}
(vm)$ gsutil cp ./data/dir/coco/raw-data/annotations/*.json ${DATA_DIR}