本教程介绍如何在 Cloud TPU 和 GKE 上训练 Mask RCNN 模型。
目标
- 创建 Cloud Storage 存储分区以保存数据集和模型输出。
- 创建一个 GKE 集群来管理您的 Cloud TPU 资源。
- 下载 Kubernetes 作业规范,其中指明了使用 TensorFlow 在 Cloud TPU 上训练 Mask RCNN 所需的资源。
- 在 GKE 集群中运行作业以开始训练模型。
- 检查日志和模型输出。
费用
本教程使用 Google Cloud 的以下收费组件:- Compute Engine
- Cloud TPU
- Cloud Storage
请使用价格计算器根据您的预计使用情况来估算费用。 Google Cloud 新用户可能有资格申请免费试用。
准备工作
-
登录您的 Google 帐号。
如果您还没有 Google 帐号,请注册新帐号。
-
在 Google Cloud Console 的项目选择器页面上,选择或创建一个 Google Cloud 项目。
-
确保您的 Cloud 项目已启用结算功能。 了解如何确认您的项目是否已启用结算功能。
-
通过在 Cloud Console 中打开以下链接,启用本教程所需的 API:
结合使用 GKE 和 Cloud TPU 时,您的项目会 使用 Google Cloud 的计费组件。如需估算您的费用,请查看 Cloud TPU 价格和 GKE 价格;并且在使用完资源之后,按照相关说明清理资源。
创建 Cloud Storage 存储分区
您需要一个 Cloud Storage 存储分区来存储模型的训练数据和训练结果。本教程中使用的 ctpu up
工具会为 Cloud TPU 服务帐号设置默认权限。如果您需要更精细的权限,请查看访问级层权限。
转到 Cloud Console 上的 Cloud Storage 页面。
创建一个新的存储分区,并指定以下选项:
- 您选择的唯一名称。
- 默认存储类别:
Regional
- 选择数据存储位置:
region
位置:us-central1
- 选择数据的默认存储类别:标准
- 选择如何控制对象的访问权限:设置对象级层和存储分区级层权限
在 GKE 上创建集群
按照以下说明操作,使用 gcloud
命令行工具设置环境并创建一个支持 Cloud TPU 的 GKE 集群。
打开一个 Cloud Shell 窗口。
配置
gcloud
命令行工具,以使用要在其中创建 Cloud TPU 的项目。gcloud config set project ${PROJECT_ID}
当您第一次在新的 Cloud Shell 虚拟机中运行此命令时,系统会显示
Authorize Cloud Shell
页面。点击页面底部的Authorize
以允许gcloud
使用您的凭据进行 GCP API 调用。为 Cloud TPU 项目创建服务帐号。
gcloud beta services identity create --service tpu.googleapis.com --project $PROJECT_ID
该命令将返回以下格式的 Cloud TPU 服务帐号:
service-PROJECT_NUMBER@cloud-tpu.iam.gserviceaccount.com
授权 Cloud TPU 访问 Cloud Storage 存储分区
您需要向 Cloud TPU 提供对 Cloud Storage 对象的读取/写入权限。为此,您必须向 Cloud TPU 使用的服务帐号授予必需的访问权限。请按照相关指南授予对存储分区的访问权限。
指定要在其中使用 Cloud TPU 资源的地区。 本示例使用的是
us-central1-b
地区:$ gcloud config set compute/zone us-central1-b
使用
gcloud container clusters
命令在 GKE 上创建一个支持 Cloud TPU 的集群。请注意,必须在提供 Cloud TPU 的地区创建 GKE 集群及其节点池,如上文中关于环境变量的部分所述。以下命令会创建一个名为tpu-models-cluster
的集群:$ gcloud container clusters create tpu-models-cluster \ --cluster-version=1.13 \ --scopes=cloud-platform \ --enable-ip-alias \ --enable-tpu \ --machine-type=n1-standard-4
在上述命令中:
--cluster-version=1.13
表示该集群将使用最新的 Kubernetes 1.16 版本。您必须使用1.13.4-gke.5
或更高版本。--scopes=cloud-platform
可确保集群中的所有节点都可以访问上面定义为project-name 的 Google Cloud 中的 Cloud Storage 存储分区。该集群和存储分区必须位于同一项目中。请注意,Pod 默认继承其所部署到的节点的访问权限范围。此标志向在该集群中运行的所有 Pod 授予cloud-platform
范围。如果您想要根据每个 Pod 来限制访问权限,请参阅 GKE 指南通过服务帐号进行身份验证。--enable-ip-alias
表示集群使用别名 IP 范围。如需在 GKE 上使用 Cloud TPU,这是必需的设置。--enable-tpu
表示集群必须支持 Cloud TPU。- 运行此作业需要
machine-type=n1-standard-4
。
命令运行完毕时,将显示如下确认消息:
kubeconfig entry generated for tpu-models-cluster. NAME LOCATION MASTER_VERSION MASTER_IP MACHINE_TYPE NODE_VERSION NUM_NODES STATUS tpu-models-cluster us-central1-b 1.13.6-gke.5 35.232.204.86 n1-standard-4 1.13.6-gke.5 3 RUNNING
处理训练数据
您的下一个任务是创建一个作业。在 Google Kubernetes Engine 中,作业是一个表示有限任务的控制器对象。您要创建的第一个作业将下载和处理用于训练 Mask RCNN 模型的 COCO 数据集。
在 shell 环境中,创建名为
download_and_preprocess_coco_k8s.yaml
的文件,如下所示。或者,您可以从 GitHub 下载此文件。找到名为
DATA_BUCKET
的环境变量。将此变量的value
更新为您要在 Cloud Storage 存储分区上存储 COCO 数据集的路径。例如:“gs://my-maskrcnn-bucket/coco/”运行以下命令以在集群中创建作业:
kubectl create -f download_and_preprocess_coco_k8s.yaml
作业可能需要几分钟才能启动。您可以使用以下命令查看集群中运行的作业的状态:
kubectl get pods -w
此命令返回有关集群中运行的所有作业的信息。例如:
NAME READY STATUS RESTARTS AGE download-and-preprocess-coco-xmx9s 0/1 Pending 0 10s
-w
标志指示该命令监视 pod 状态的任何更改。几分钟后,您的作业状态应更新为 Running
。
如果您的作业遇到任何问题,可以将其删除,然后重试。为此,请先运行 kubectl delete -f download_and_preprocess_coco_k8s.yaml
命令,然后再次运行 kubectl create
命令。
下载和预处理完成后,您就可以运行该模型了。
运行 Mask RCNN 模型
现已全部准备就绪,您可使用 Cloud TPU 和 GKE 来运行 Mask RCNN 模型了。
在 shell 环境中,创建名为
mask_rcnn_k8s.yaml
的文件,如下所示。或者,您可以从 GitHub 下载此文件。找到名为
DATA_BUCKET
的环境变量。将此变量的value
更新为您在download_and_preprocess_coco_k8s.yaml
脚本中指定的DATA_BUCKET
文件的路径。找到名为
MODEL_BUCKET
的环境变量。将此变量的value
更新为 Cloud Storage 存储分区的路径。此位置就是脚本保存训练输出的位置。如果指定的MODEL_BUCKET
不存在,则会在您运行作业时创建。运行以下命令以在集群中创建作业:
kubectl create -f mask_rcnn_k8s.yaml
运行此命令时,将显示以下确认消息:
job "mask-rcnn-gke-tpu" created
如处理训练数据部分所述,该作业可能需要几分钟才能启动。使用以下命令查看在您的集群中运行的作业的状态:
kubectl get pods -w
查看训练状态
您可以使用 kubectl
命令行实用程序跟踪训练状态。
运行以下命令以获取作业的状态:
kubectl get pods
在训练过程中,pod 的状态应为
RUNNING
。如需获取有关训练过程的其他信息,请运行以下命令:
$ kubectl logs job name
其中
job name
是作业的完整名称,例如mask-rcnn-gke-tpu-abcd
。您还可以在 Cloud Console 上的 GKE 工作负载信息中心检查输出。
请注意,第一个条目需要一段时间后才会在日志中显示。您可能会看到如下内容:
I0622 18:14:31.617954 140178400511808 tf_logging.py:116] Calling model_fn. I0622 18:14:40.449557 140178400511808 tf_logging.py:116] Create CheckpointSaverHook. I0622 18:14:40.697138 140178400511808 tf_logging.py:116] Done calling model_fn. I0622 18:14:44.004508 140178400511808 tf_logging.py:116] TPU job name worker I0622 18:14:45.254548 140178400511808 tf_logging.py:116] Graph was finalized. I0622 18:14:48.346483 140178400511808 tf_logging.py:116] Running local_init_op. I0622 18:14:48.506665 140178400511808 tf_logging.py:116] Done running local_init_op. I0622 18:14:49.135080 140178400511808 tf_logging.py:116] Init TPU system I0622 18:15:00.188153 140178400511808 tf_logging.py:116] Start infeed thread controller I0622 18:15:00.188635 140177578452736 tf_logging.py:116] Starting infeed thread controller. I0622 18:15:00.188838 140178400511808 tf_logging.py:116] Start outfeed thread controller I0622 18:15:00.189151 140177570060032 tf_logging.py:116] Starting outfeed thread controller. I0622 18:15:07.316534 140178400511808 tf_logging.py:116] Enqueue next (100) batch(es) of data to infeed. I0622 18:15:07.316904 140178400511808 tf_logging.py:116] Dequeue next (100) batch(es) of data from outfeed. I0622 18:16:13.881397 140178400511808 tf_logging.py:116] Saving checkpoints for 100 into gs://<bucket-name>/mask_rcnn/model.ckpt. I0622 18:16:21.147114 140178400511808 tf_logging.py:116] loss = 1.589756, step = 0 I0622 18:16:21.148168 140178400511808 tf_logging.py:116] loss = 1.589756, step = 0 I0622 18:16:21.150870 140178400511808 tf_logging.py:116] Enqueue next (100) batch(es) of data to infeed. I0622 18:16:21.151168 140178400511808 tf_logging.py:116] Dequeue next (100) batch(es) of data from outfeed. I0622 18:17:00.739207 140178400511808 tf_logging.py:116] Enqueue next (100) batch(es) of data to infeed. I0622 18:17:00.739809 140178400511808 tf_logging.py:116] Dequeue next (100) batch(es) of data from outfeed. I0622 18:17:36.598773 140178400511808 tf_logging.py:116] global_step/sec: 2.65061 I0622 18:17:37.040504 140178400511808 tf_logging.py:116] examples/sec: 2698.56 I0622 18:17:37.041333 140178400511808 tf_logging.py:116] loss = 2.63023, step = 200 (75.893 sec)
在
gs://<bucket-name>/mask-rcnn/model.ckpt
中查看经过训练的模型。gsutil ls -r gs://bucket-name/mask-rcnn/model.ckpt
清理
在 GKE 上使用 Cloud TPU 后,请清理资源,以避免您的 Google Cloud 帐号产生额外扣款。
如果您尚未为此会话设置项目和地区,请立即设置。请参阅本指南前面的说明。然后按照下列清理步骤进行操作:
运行以下命令以删除 GKE 集群
tpu-models-cluster
,将 project-name 替换为您的 Google Cloud 项目名称:$ gcloud container clusters delete tpu-models-cluster --project=project-name
检查完数据后,使用
gsutil
命令删除您在本教程期间创建的 Cloud Storage 存储分区。将 bucket-name 替换为您的 Cloud Storage 存储分区的名称:$ gsutil rm -r gs://bucket-name
请参阅 Cloud Storage 价格指南,了解免费存储限制和其他价格信息。
后续步骤
- 探索 TensorBoard 中的 TPU 工具。
使用以下作业规范之一运行更多模型和数据集检索作业:
- 在 GKE 上下载并预处理 COCO 数据集。
- 在 GKE 上下载并预处理 ImageNet。
- 使用 Cloud TPU 和 GKE 训练 AmoebaNet-D。
- 使用 Cloud TPU 和 GKE 训练 Inception v3。
- 使用 Cloud TPU 和 GKE 训练 RetinaNet。
使用更多 TPU 示例进行实验。