在 Cloud TPU 上运行 MNIST (TF 2.x)

本教程包含 MNIST 模型的简要说明、有关下载 MNIST TensorFlow TPU 代码示例的说明,以及在 Cloud TPU 上运行代码的指南。

免责声明

本教程使用第三方数据集。Google 不对此数据集的有效性或任何其他方面提供任何表示、担保或其他保证。

模型说明

MNIST 数据集包含大量 0 到 9 之间的手写数字图片,以及标识每个图片中数字的标签。

本教程将训练一个机器学习模型作为示例,以根据 MNIST 数据集对图片进行分类。经过训练之后,该模型根据从 MNIST 数据集中获取的手写图像而学到的内容,将输入图像分为 10 个类别(0 到 9)。然后,您可以向模型发送它之前未曾见过的图像,之后模型会根据其在训练期间学到的内容来识别图像中的数字。

MNIST 数据集分为以下三个部分:

  • 60000 个训练数据示例
  • 10000 个测试数据示例
  • 5000 个验证数据示例

此模型包含 7 个层:

  • 2 个卷积
  • 2 个池化
  • 2 个密集(全连接)
  • 1 个漏失

损失通过分类交叉熵计算得出。

此版本的 MNIST 模型使用 Keras API,这是在 Cloud TPU 上构建和运行机器学习模型的推荐方法。

Keras 通过隐藏大部分低级实现来简化模型开发流程,从而让您更轻松在 TPU 和其他测试平台(例如 GPU 或 CPU)之间切换。

目标

  • 创建 Cloud Storage 存储分区以保存数据集和模型输出。
  • 运行训练作业。
  • 验证输出结果。

费用

本教程使用 Google Cloud 的以下收费组件:

  • Compute Engine
  • Cloud TPU
  • Cloud Storage

请使用价格计算器根据您的预计使用情况来估算费用。 Google Cloud 新用户可能有资格申请免费试用

准备工作

本部分介绍如何设置 Cloud Storage 存储分区和 Compute Engine 虚拟机。

  1. 打开一个 Cloud Shell 窗口。

    打开 Cloud Shell

  2. 为项目 ID 创建一个变量。

    export PROJECT_ID=project-id
    
  3. 配置 gcloud 命令行工具,以使用要在其中创建 Cloud TPU 的项目。

    gcloud config set project ${PROJECT_ID}
    

    当您第一次在新的 Cloud Shell 虚拟机中运行此命令时,系统会显示 Authorize Cloud Shell 页面。点击页面底部的 Authorize 以允许 gcloud 使用您的凭据进行 GCP API 调用。

  4. 为 Cloud TPU 项目创建服务帐号。

    gcloud beta services identity create --service tpu.googleapis.com --project $PROJECT_ID
    

    该命令将返回以下格式的 Cloud TPU 服务帐号:

    service-PROJECT_NUMBER@cloud-tpu.iam.gserviceaccount.com
    

  5. 使用以下命令创建 Cloud Storage 存储分区:

    gsutil mb -p ${PROJECT_ID} -c standard -l us-central1 -b on gs://bucket-name
    

    此 Cloud Storage 存储分区存储您用于训练模型的数据和训练结果。本教程中使用的 gcloud 命令会为您在上一步中设置的 Cloud TPU 服务帐号设置默认权限。如果您需要更精细的权限,请查看访问级层权限

  6. 使用 gcloud 命令启动 Compute Engine 虚拟机和 Cloud TPU。 您使用的命令取决于您使用的是 TPU 虚拟机还是 TPU 节点。如需了解详情,请参阅系统架构

    TPU 虚拟机

    $ gcloud alpha compute tpus tpu-vm create mnist-tutorial \
    --zone=us-central1-b \
    --accelerator-type=v3-8 \
    --version=v2-alpha
    

    命令标志说明

    zone
    拟在其中创建 Cloud TPU 的可用区
    accelerator-type
    要创建的 Cloud TPU 的类型
    version
    Cloud TPU 运行时版本。

    TPU 节点

    $ gcloud compute tpus execution-groups create \
    --name=mnist-tutorial \
    --zone=us-central1-b \
    --tf-version=2.4.1 \
    --machine-type=n1-standard-1 \
    --accelerator-type=v3-8
    

    命令标志说明

    name
    要创建的 Cloud TPU 的名称。
    zone
    拟在其中创建 Cloud TPU 的可用区
    tf-version
    gcloud 命令在您的虚拟机上安装的 Tensorflow 版本。
    machine-type
    要创建的 Compute Engine 虚拟机的机器类型
    accelerator-type
    要创建的 Cloud TPU 的类型

    如需详细了解 gcloud 命令,请参阅 gcloud 参考文档

  7. gcloud compute tpus execution-groups 命令执行完毕后,验证 shell 提示符已从 username@projectname 更改为 username@vm-name。此变化表明您现已登录 Compute Engine 虚拟机。

    如果您未连接到 Compute Engine 实例,可以运行以下命令进行连接:

    TPU 虚拟机

    gcloud alpha compute tpus tpu-vm ssh mnist-tutorial --zone=us-central1-b
    

    TPU 节点

    gcloud compute ssh mnist-tutorial --zone=us-central1-b
    

    在您继续按照这些说明操作时,请在虚拟机会话窗口中运行以 (vm)$ 开头的每个命令。

  8. 为 TPU 名称创建环境变量。

    TPU 虚拟机

    (vm)$ export TPU_NAME=local
    

    TPU 节点

    (vm)$ export TPU_NAME=mnist-tutorial
    
  9. 安装 TensorFlow 要求。

    您使用的命令取决于您使用的是 TPU 虚拟机还是 TPU 节点。

    TPU 虚拟机

    (vm)$ git clone https://github.com/tensorflow/models.git
    (vm)$ pip3 install -r models/official/requirements.txt
    

    TPU 节点

    (vm)$ sudo pip3 install tensorflow-model-optimization>=0.1.3
    

单个 Cloud TPU 设备训练

MNIST TPU 模型的源代码可在 GitHub 上找到。

  1. 设置以下变量。将 bucket-name 替换为您的存储分区名称。

    (vm)$ export STORAGE_BUCKET=gs://bucket-name
    (vm)$ export MODEL_DIR=${STORAGE_BUCKET}/mnist
    (vm)$ export DATA_DIR=${STORAGE_BUCKET}/data
    
  2. 设置 PYTHONPATH 环境变量。

    TPU 虚拟机

    (vm)$ export PYTHONPATH="${PWD}/models:${PYTHONPATH}"
    

    TPU 节点

    (vm)$ export PYTHONPATH="${PYTHONPATH}:/usr/share/models"
    
  3. 切换至存储模型的目录。

    TPU 虚拟机

    (vm)$ cd ~/models/official/vision/image_classification
    

    TPU 节点

    (vm)$ cd /usr/share/models/official/vision/image_classification
    
  4. 运行 MNIST 训练脚本:

    (vm)$ python3 mnist_main.py \
      --tpu=${TPU_NAME} \
      --model_dir=${MODEL_DIR} \
      --data_dir=${DATA_DIR} \
      --train_epochs=10 \
      --distribution_strategy=tpu \
      --download
    

    命令标志说明

    tpu
    Cloud TPU 的名称。如果您在设置 Compute Engine 虚拟机和 Cloud TPU 时未指定此名称,系统会默认使用您的用户名。
    model_dir
    在训练期间存储检查点和摘要的 Cloud Storage 存储分区。您可以使用现有文件夹加载根据相同大小和 TensorFlow 版本的先前生成 TPU 创建的检查点。
    data_dir
    训练输入的 Cloud Storage 路径。在此示例中,该路径设置为 fake_imagenet 数据集。
    train_epochs
    训练模型的周期数。
    distribution_strategy
    如需在 Cloud TPU 上训练 ResNet 模型,请将 distribution_strategy 设置为 tpu
    download
    如果设置为 true,该脚本会下载并预处理 MNIST 数据集(如果该数据集尚未下载)。

训练脚本在 v3-8 Cloud TPU 上运行的时间不到 5 分钟,并显示类似以下内容的输出:

Run stats:
{
  'accuracy_top_1': 0.9762369990348816,
  'eval_loss': 0.07863274961709976,
  'loss': 0.1111728847026825,
  'training_accuracy_top_1': 0.966645359992981
}

清理

为避免因本教程中使用的资源导致您的 Google Cloud 帐号产生费用,请删除包含这些资源的项目,或者保留项目但删除各个资源。

  1. 断开与 Compute Engine 实例的连接(如果您尚未这样做):

    (vm)$ exit
    

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

  2. 删除您的 Cloud TPU 和 Compute Engine 资源。 您用于删除资源的命令取决于您使用的是 TPU 虚拟机还是 TPU 节点。如需了解详情,请参阅系统架构

    TPU 虚拟机

    $ gcloud alpha compute tpus tpu-vm delete mnist-tutorial \
    --zone=us-central1-b
    

    TPU 节点

    $ gcloud compute tpus execution-groups delete mnist-tutorial \
    --zone=us-central1-b
    
  3. 通过运行 gcloud compute tpus execution-groups list 验证资源是否已删除。删除操作可能需要几分钟时间才能完成。如下所示的响应表明实例已成功删除。

    $ gcloud compute tpus execution-groups list --zone=us-central1-b
    
    NAME             STATUS
    
  4. 使用 gsutil 删除 Cloud Storage 存储分区,如下所示。将 bucket-name 替换为您的 Cloud Storage 存储分区的名称。

    $ gsutil rm -r gs://bucket-name
    

后续步骤

在本教程中,您已使用示例数据集训练 minst 模型。此训练的结果(在大多数情况下)不能用于推断。要使用模型进行推断,您可以在公开提供的数据集或您自己的数据集上训练数据。在 Cloud TPU 上训练的模型要求数据集采用 TFRecord 格式。

您可以使用数据集转换工具示例将图片分类数据集转换为 TFRecord 格式。如果您未使用图片分类模型,则必须自行将数据集转换为 TFRecord 格式。如需了解详情,请参阅 TFRecord 和 tf.Example

超参数调节

如需使用数据集提升模型的性能,您可以调节模型的超参数。您可以在 GitHub 上寻找所有 TPU 支持模型通用的超参数的相关信息。您可以在每个模型的源代码中寻找模型专用超参数的相关信息。如需详细了解超参数调节,请参阅超参数调节概览使用超参数调节服务调节超参数

推理

训练模型后,您可以使用该模型进行推断(也称为预测)。AI Platform 是一款基于云的解决方案,用于开发、训练部署机器学习模型。部署模型后,您可以使用 AI Platform Prediction 服务