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

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

  • Mask-RCNN
  • Retinanet
  • ShapeMask

在 Cloud TPU 上训练模型之前,您必须先准备训练数据。

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

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

要完全下载、预处理 COCO 数据集并将其上传到 Cloud Storage 存储桶,大约需要 2 小时。

  1. Cloud Shell 中,使用您的项目 ID 配置 gcloud

    export PROJECT_ID=project-id
    gcloud config set project ${PROJECT_ID}
  2. Cloud Shell 中,使用以下命令创建 Cloud Storage 存储桶:

    gcloud storage buckets create gs://bucket-name --project=${PROJECT_ID} --location=us-central2
  3. 创建一个 Compute Engine 虚拟机,以下载并预处理数据集。如需了解详情,请参阅创建和启动 Compute Engine 实例

    $ gcloud compute instances create vm-name \
        --zone=us-central2-b \
        --image-family=ubuntu-2204-lts \
        --image-project=ubuntu-os-cloud \
        --machine-type=n1-standard-16 \
        --boot-disk-size=300GB
  4. 使用 SSH 连接到 Compute Engine 虚拟机:

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

    连接到虚拟机后,shell 提示符会从 username@projectname 更改为 username@vm-name

  5. 设置两个变量,一个用于先前创建的存储桶,另一个用于保存存储桶中的训练数据 (DATA_DIR) 的目录。

    (vm)$ export STORAGE_BUCKET=gs://bucket-name
    (vm)$ export DATA_DIR=${STORAGE_BUCKET}/coco
  6. 安装预处理数据所需的软件包。

    (vm)$ sudo apt-get update && \
      sudo apt-get install python3-pip && \
      sudo apt-get install -y python3-tk && \
      pip3 install --user Cython matplotlib opencv-python-headless pyyaml Pillow numpy absl-py tensorflow && \
      pip3 install --user "git+https://github.com/cocodataset/cocoapi#egg=pycocotools&subdirectory=PythonAPI"
  7. 运行 download_and_preprocess_coco.sh 脚本,将 COCO 数据集转换为训练应用所需的一组 TFRecord 文件 (*.tfrecord)。

    (vm)$ git clone https://github.com/tensorflow/tpu.git
    (vm)$ sudo bash tpu/tools/datasets/download_and_preprocess_coco.sh ./data/dir/coco

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

  8. 将数据复制到 Cloud Storage 存储桶。

    将数据转换为 TFRecord 格式后,使用 gcloud CLI 将数据从本地存储空间复制到 Cloud Storage 存储桶。您还必须复制注解文件。这些文件有助于验证模型的性能。

    (vm)$ gcloud storage cp ./data/dir/coco/*.tfrecord ${DATA_DIR}
    (vm)$ gcloud storage cp ./data/dir/coco/raw-data/annotations/*.json ${DATA_DIR}
  9. 与 Compute Engine 虚拟机断开连接:

    (vm)$ exit

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

  10. 删除您的 Compute Engine 虚拟机:

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