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

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

免责声明

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

模型说明

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

本教程将训练一个机器学习模型，以根据 MNIST 数据集。训练完成后，模型将传入的图片分为 10 张 分类（0 到 9）。 MNIST 数据集。然后，您可以向模型发送它之前未曾见过的图像，之后模型会根据其在训练期间学到的内容来识别图像中的数字。

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. 配置 Google Cloud CLI 以使用要在其中创建项目的项目 Cloud TPU。

    gcloud config set project ${PROJECT_ID}

    当您第一次在新的 Cloud Shell 虚拟机中运行此命令时，系统会显示 Authorize Cloud Shell 页面。点击页面底部的 Authorize 以允许 gcloud 使用您的凭据进行 Google Cloud 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 存储分区：

    gcloud storage buckets create gs://bucket-name \
--project=${PROJECT_ID} \
--location=us-central1

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

  6. 使用 gcloud 命令创建 TPU 虚拟机。

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

    命令标志说明

    zone
    拟在其中创建 Cloud TPU 的可用区。

    accelerator-type
    加速器类型用于指定您要创建的 Cloud TPU 的版本和大小。 如需详细了解每个 TPU 版本支持的加速器类型，请参阅 TPU 版本。

    version
    Cloud TPU 软件版本。

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

  7. 使用 SSH 连接到 Compute Engine 实例。连接到网络后 您的 Shell 提示符会从 username@projectname 更改为 username@vm-name:

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

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

    (vm)$ export TPU_NAME=local

  9. 安装 TensorFlow 要求。

    (vm)$ pip3 install -r /usr/share/tpu/models/official/requirements.txt

训练模型

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 环境变量。

    (vm)$ export PYTHONPATH="${PYTHONPATH}:/usr/share/tpu/models"

  3. 创建 TPU 时，如果您将 --version 参数设置为以 -pjrt 时，请设置以下环境变量以启用 PJRT 运行时： 

      (vm)$ export NEXT_PLUGGABLE_DEVICE_USE_C_API=true
  (vm)$ export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so

  4. 切换至存储模型的目录。

    (vm)$ cd /usr/share/tpu/models/official/legacy/image_classification

  5. 运行 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 存储分区。您可以使用现有文件夹加载之前的内容 在相同大小的 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 资源。

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

  3. 通过运行 gcloud compute tpus tpu-vm list 验证资源是否已删除。删除操作可能需要几分钟时间才能完成。以下命令的输出 不应包含本教程中创建的任何资源：

    $ gcloud compute tpus tpu-vm list --zone=us-central1-b

  4. 使用以下命令删除 Cloud Storage 存储桶。将 bucket-name 替换为您的 Cloud Storage 存储分区的名称。

    $ gcloud storage rm gs://bucket-name --recursive

后续步骤

TensorFlow Cloud TPU 教程通常使用示例数据集来训练模型。此训练的结果不能用于推断。接收者 使用模型进行推理，就可以使用公开可用的 或您自己的数据集。在 Cloud TPU 上训练的 TensorFlow 模型 通常需要将数据集 TFRecord 格式。

您可以使用数据集转换工具示例将图片分类数据集转换为 TFRecord 格式。如果您使用的不是图片 分类模型，您需要将数据集转换为 TFRecord 格式 。如需了解详情，请参阅 TFRecord 和 tf.Example。

超参数调节

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

推断

训练模型后，即可将其用于推理（也称为 预测）。您可以使用 Cloud TPU 推理转换器工具准备和优化 TensorFlow 模型，以便在 Cloud TPU v5e 上进行推理。如需详细了解 Cloud TPU v5e 上的推理，请参阅 Cloud TPU v5e 推理简介。