下载、预处理和上传 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 数据集并将其上传到 Google Cloud 存储桶,大约需要 2 个小时。

  1. Cloud Shell 中,使用以下命令创建 Cloud Storage 存储桶:

    gsutil mb -p ${PROJECT_ID} -c standard -l europe-west4 gs://bucket-name
    
  2. 启动 Compute Engine 虚拟机实例。

    此虚拟机实例将仅用于下载和预处理 COCO 数据集。在 instance-name 中填写您选择的名称。

    $ gcloud compute tpus execution-groups create \
     --vm-only \
     --name=instance-name \
     --zone=europe-west4-a \
     --disk-size=300 \
     --machine-type=n1-standard-16 \
     --tf-version=2.6.0
    

    命令标志说明

    vm-only
    仅创建虚拟机。默认情况下,gcloud compute tpus execution-groups 命令会同时创建虚拟机和 Cloud TPU。
    name
    要创建的 Cloud TPU 的名称。
    zone
    拟在其中创建 Cloud TPU 的区域
    disk-size
    gcloud compute tpus execution-groups 命令创建的虚拟机的硬盘大小(以 GB 为单位)。
    machine-type
    要创建的 Compute Engine 虚拟机的机器类型
    tf-version
    在虚拟机上安装的 Tensorflow gcloud compute tpus execution-groups 的版本。
  3. 如果您未自动登录 Compute Engine 实例,请通过运行以下 ssh 命令进行登录。登录虚拟机后,shell 提示符会从 username@projectname 更改为 username@vm-name

      $ gcloud compute ssh instance-name --zone=europe-west4-a
      

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

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

    (vm)$ sudo apt-get install -y python3-tk && \
      pip3 install --user Cython matplotlib opencv-python-headless pyyaml Pillow && \
      pip3 install --user "git+https://github.com/cocodataset/cocoapi#egg=pycocotools&subdirectory=PythonAPI"
    
  6. 运行 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 小时才能完成。

  7. 将数据复制到 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}
    
  8. 清理虚拟机资源

    将 COCO 数据集转换为 TFRecord 并将其复制到 Cloud Storage 存储桶的 DATA_DIR 后,您便可以删除 Compute Engine 实例。

    与 Compute Engine 实例断开连接:

    (vm)$ exit
    

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

  9. 删除您的 Compute Engine 实例。

      $ gcloud compute instances delete instance-name
        --zone=europe-west4-a