在 Cloud TPU 上训练 (TF 2.x)

设置 Cloud Storage 存储桶变量

设置以下环境变量，将 bucket-name 替换为 Cloud Storage 存储桶的名称：

(vm)$ export STORAGE_BUCKET=gs://bucket-name

export MODEL_DIR=${STORAGE_BUCKET}/resnet-rs-2x
export IMAGENET_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
export PYTHONPATH=/usr/share/models
export TPU_NAME=resnet-rs-tutorial

训练应用应该能够访问 Cloud Storage 中的训练数据。在训练期间，训练应用还会使用您的 Cloud Storage 存储桶来存储检查点。

Cloud TPU 单设备训练和评估

ImageNet 是一个图片数据库。数据库中的图片被整理为一个层次结构，该层次结构中的每个节点由成百上千个图片进行描述。

本教程使用演示版的完整 ImageNet 数据集，该数据集又称为 fake_imagenet fake_imagenet。此演示版本可用于测试教程，同时降低通常与使用完整 ImageNet 数据库运行模型相关的存储和时间要求。

此 fake_imagenet 数据集位于 Cloud Storage 上的以下位置：

gs://cloud-tpu-test-datasets/fake_imagenet

此 fake_imagenet 数据集仅用于了解如何使用 Cloud TPU 并验证端到端性能。准确率数字和保存的模型并无实际意义。

如需了解如何下载和处理完整 ImageNet 数据集，请参阅下载、预处理和上传 ImageNet 数据集。

使用 gcloud 命令创建 Cloud TPU。

gcloud compute tpus execution-groups create \
 --tpu-only \
 --accelerator-type=v3-8  \
 --name=resnet-rs-tutorial \
 --zone=europe-west4-a \
 --tf-version=2.12.0

设置 TPU_NAME 名称变量。

(vm)$ export TPU_NAME=resnet-rs-tutorial

运行训练脚本。

python3 /usr/share/models/official/vision/beta/train.py \
--experiment=resnet_rs_imagenet \
--mode=train_and_eval \
--model_dir=$MODEL_DIR \
--tpu=$TPU_NAME \
--config_file=/usr/share/models/official/vision/beta/configs/experiments/image_classification/imagenet_resnetrs50_i160.yaml \
--params_override="task.train_data.input_path=$IMAGENET_DIR/train*, task.validation_data.input_path=$IMAGENET_DIR/valid*, trainer.train_steps=100"

这将训练 ResNet-RS 进行 100 个训练步骤，并在 5 分钟之内完成 v3-8 TPU 节点。训练脚本输出应包括如下所示的文本：

{
  'train_loss': 1.435225,
  'train_accuracy': 0.00084427913
}

训练脚本也会执行评估。评估输出应包含如下文本：

Run stats:
{
  'eval_loss': 0.861013,
  'eval_acc': 0.001,
  'train_loss': 1.435225,
  'train_acc': 0.00084427913,
  'step_timestamp_log': [
    'BatchTimestamp<batch_index: 0,
    timestamp: 1606330585.7613473>',
    'BatchTimestamp<batch_index: 500,
    timestamp: 1606330883.8486104>',
    'BatchTimestamp<batch_index: 1000,
    timestamp: 1606331119.515312>',
    'BatchTimestamp<batch_index: 1251,
    timestamp: 1606331240.7516596>'
  ],
  'train_finish_time': 1606331296.395158,
  'avg_exp_per_second': 1951.6983246161021
}

如需将 ResNet-RS 模型收敛，请省略 trainer.train_steps=100 参数，如以下脚本所示。训练和评估是一起完成的。

python3 /usr/share/models/official/vision/beta/train.py \
  --experiment=resnet_rs_imagenet \
  --mode=train_and_eval \
  --model_dir=$MODEL_DIR \
  --tpu=$TPU_NAME \
  --config_file=/usr/share/models/official/vision/beta/configs/experiments/image_classification/imagenet_resnetrs50_i160.yaml \
  --params_override="task.train_data.input_path=$IMAGENET_DIR/train*, task.validation_data.input_path=$IMAGENET_DIR/valid*"

由于训练和评估是利用 fake_imagenet 数据集完成的，因此输出结果不会反映实际输出，实际输出在利用实际数据集执行训练和评估时才会出现。

此时，您可以结束本教程并清理 Google Cloud 资源，也可以进一步了解如何在 Cloud TPU Pod 上运行模型。

使用较大的模型

ResNet-RS 提供不同大小的模型系列，模型越大，通常代价越高，计算费用越高。如需了解详情，请参阅重新访问 ResNet：改进的训练和扩缩策略。

您可以通过更改以下命令中的 config_file 来选择要训练的模型的大小。

python3 /usr/share/models/official/vision/beta/train.py \
  --experiment=resnet_rs_imagenet \
  --mode=train_and_eval \
  --model_dir=$MODEL_DIR \
  --tpu=$TPU_NAME \
  --config_file=/usr/share/models/official/vision/beta/configs/experiments/image_classification/imagenet_resnetrs200_i256.yaml \
  --params_override="task.train_data.input_path=$IMAGENET_DIR/train*, task.validation_data.input_path=$IMAGENET_DIR/valid*"

可用配置位于虚拟机的 /usr/share/models/official/vision/beta/configs/experiments/ 中。

使用 Cloud TPU Pod 扩缩模型

在 Cloud TPU Pod 上训练模型可能需要对训练脚本进行一些更改。如需了解详情，请参阅在 TPU Pod 上训练。

您可以使用 Cloud TPU Pod 扩缩模型，以便更快获得结果。完全受支持的 ResNet-RS-50 模型可与以下 Pod 切片配合使用：

v2-32
v3-32

借助 Cloud TPU Pod，训练和评估可同时完成。

使用 Cloud TPU Pod 进行训练

删除为在单台设备上训练模型而创建的 Cloud TPU 资源。

gcloud compute tpus execution-groups delete resnet-rs-tutorial \
  --zone=europe-west4-a \
  --tpu-only

为 Pod 创建新的 Cloud TPU 资源，并使用 accelerator-type 参数指定要使用的 Pod 切片。例如，以下命令使用 v2-32 Pod 切片。

gcloud compute tpus execution-groups  create --name=resnet-rs-tutorial \
  --accelerator-type=v3-32  \
  --zone=europe-west4-a \
  --tf-version=2.12.0 \
  --tpu-only

运行训练脚本。

python3 /usr/share/models/official/vision/beta/train.py \
--experiment=resnet_rs_imagenet \
--mode=train_and_eval \
--model_dir=$MODEL_DIR \
--tpu=$TPU_NAME \
--config_file=/usr/share/models/official/vision/beta/configs/experiments/image_classification/imagenet_resnetrs50_i160.yaml \
--params_override="task.train_data.input_path=$IMAGENET_DIR/train*, task.validation_data.input_path=$IMAGENET_DIR/valid*, trainer.train_steps=100"

这将训练 ResNet-RS 进行 100 个训练步骤，并在 5 分钟之内完成 v3-8 TPU 节点。训练脚本输出应包括如下所示的文本：

{
  'train_loss': 1.435225,
  'train_accuracy': 0.00084427913
}

训练脚本也会执行评估。评估输出应包含如下文本：

Run stats:
{
  'eval_loss': 0.861013,
  'eval_acc': 0.001,
  'train_loss': 1.435225,
  'train_acc': 0.00084427913,
  'step_timestamp_log': [
    'BatchTimestamp<batch_index: 0,
    timestamp: 1606330585.7613473>',
    'BatchTimestamp<batch_index: 500,
    timestamp: 1606330883.8486104>',
    'BatchTimestamp<batch_index: 1000,
    timestamp: 1606331119.515312>',
    'BatchTimestamp<batch_index: 1251,
    timestamp: 1606331240.7516596>'
  ],
  'train_finish_time': 1606331296.395158,
  'avg_exp_per_second': 1951.6983246161021
}

训练和评估是一起完成的。每个周期都有 1251 个步骤，总计 112590 个训练步骤和 48 个评估步骤。

文章来源：google cloud

推荐阅读

更多芯擎AI开发板干货请关注芯擎AI开发板专栏。欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。

设置 Cloud Storage 存储桶变量

Cloud TPU 单设备训练和评估

使用较大的模型

使用 Cloud TPU Pod 扩缩模型

使用 Cloud TPU Pod 进行训练

推荐阅读

目录