为 TensorFlow2 引入 Model Garden

文 / 技术项目经理 Jaeyoun Kim 和软件工程师 Jing Li

我们对 Model Garden 进行更新，以期为 TensorFlow 用户提供一个集中平台，您可用来查找最先进的模型代码示例以及 TensorFlow 2 的可重用模型库。

Model Garden 展示了关于建模的最佳实践，帮助 TensorFlow 用户充分了解与利用最近推出的 TensorFlow 2 来进行研究和产品开发。同时，TensorFlow Hub 将继续发挥其作为代码库的作用，方便用户轻松搜索即时可用的预训练模型。我们计划在 TensorFlow Hub 中持续增加 Model Garden 中的最前沿模型，并将 TensorFlow Hub 页面关联至 Model Garden 中的模型实现。

我们提供了分布式训练的多个代码示例，以帮助开发者解决计算机视觉和自然语言处理中遇到的一些问题。Model Garden 官方代码库中的代码示例将由 TensorFlow 团队进行维护，并保证与最新 TensorFlow 2 API 保持同步。

可使用 PIP (pip install tf-models-nightly) 轻松安装 Model Garden。TensorFlow 2 用户可立即开始使用代码示例，以学习在 GPU 和 TPU 上训练模型的最佳实践。

Model Garden 中的许多模型可采用分布式方法来训练。在 TensorFlow 2 中，您可以使用分布策略 API 将训练工作负载分配给单主机/多加速器配置，以及多主机/多加速器配置。

我们将在本文中介绍一些用于构建模型的常见分布策略。查看 TensorFlow 2 可用的分布策略列表，请访问 TensorFlow.google.cn。

• MirroredStrategy — 用于多个 GPU（如找不到 GPU，则使用 CPU）
• MultiWorkerMirroredStrategy — 用于具有多个 GPU 的多个主机

• TPUStrategy — 用于 TPU 或多个 TPU 主机

在以下的两个部分（计算机视觉和自然语言处理），我们将通过示例演示，介绍如何使用分布式训练。

计算机视觉

使用 ResNet 进行图像分类

ResNet 模型可将图像中的主要对象分类为 1,000 个对象类别（例如，汽车、足球、台灯等）。Model Garden 提供了一个示例，演示如何在 GPU 和 TPU 上训练 ResNet。此模型在纯 TensorFlow 2 API 环境下编写而成，对用户友好，并使用面向对象的风格。

在多个 GPU 上进行分布式训练

您可以使用 tf.distribute.MirroredStrategy API 在多个 GPU 上训练模型。以下示例说明如何在两个 GPU 上使用 ImageNet 数据训练模型。classifier_trainer.py 是一个全新的统一框架，用于通过面向构建和训练深度学习模型的 TensorFlow 高阶 API 来训练图像分类模型（Keras compile 和 fit 方法）。

$ python3 classifier_trainer.py \
    --mode=train_and_eval \
    --model_type=resnet \
    --dataset=imagenet \
    --model_dir=$MODEL_DIR \
    --data_dir=$DATA_DIR \
    --config_file=configs/examples/resnet/imagenet/gpu.yaml \
    --params_override="runtime.num_gpus=2"

您需要指定 GPU 的数量（例如，--params_override="runtime.num_gpus=2"），以便在每个 GPU 主机的多个 GPU 上执行同步分布式训练。当您使用更多 GPU 来扩大训练规模时，还需确定批次大小、周期数和训练步数的最佳值。在配置文件中可找到这些 ImageNet 数据集相关值的示例。对于超参数调整，请参考 Keras Tuner。

在多个 GPU 主机上进行分布式训练

如要在多个 GPU 主机的多个 GPU 上进行分布式训练，则需要在每个 GPU 主机上设置TF_CONFIG环境变量，以指定构成集群的具体任务、任务地址及每个任务在集群中的角色。对于首个指定为 chief worker 的 GPU 主机的 TF_CONFIG 典型示例如下所示：

os.environ["TF_CONFIG"] = json.dumps({
    "cluster": {
        "worker": ["host1:port", "host2:port", "host3:port"]
    },
   "task": {"type": "worker", "index": 0}})

在此示例中，“worker” 部分配置三个 GPU 主机（host1、host2 和 host3），使用 MultiWorkerMirroredStrategy 执行训练。“Task” 部分指定集群中当前任务的角色。对于责任更大一些的 chief worker（例如，保存检查点），您需要将 任务类型 (task type) 设置为“Worker”，并将 任务索引 (task index) 设置为 0。MultiWorkerMirroredStrategy将自动使用每个主机上的所有可用 GPU。

在 Cloud TPU 上进行分布式训练

通过使用 TensorFlow 2，您可以轻松地在不同的硬件配置上分配和训练模型，而无需更改模型定义。用户可以使用tf.distribute.TPUStrategy在 Cloud TPU 上训练 ResNet 模型。您可以通过同一框架 (classifier_trainer.py) 来使用 Cloud TPU 训练模型。您只需为 TPU 使用不同的 YAML 配置文件并设置 --tpu=$TPU_NAME即可，其中$TPU_NAME是您的 TPU 实例在 Cloud Console 中的名称。请参阅 TPU 配置文件。

$ python3 classifier_trainer.py \
    --mode=train_and_eval \
    --model_type=resnet \
    --dataset=imagenet \
    --tpu=$TPU_NAME \
    --model_dir=$MODEL_DIR \
    --data_dir=$DATA_DIR \
    --config_file=configs/examples/resnet/imagenet/tpu.yaml

对于想使用自己的训练循环，而非使用 TensorFlow 高级 API 来构建和训练深度学习模型的用户，请查看自定义训练教程。

自然语言处理 (NLP)

今天我们介绍的第二个示例是关于 NLP 的示例。Model Garden 包含用于 BERT（Transformer 双向编码器表示）的 TensorFlow 2 实现和 BERT 变体模型（如 ALBERT）。下文我们将演示在 GPU 和 TPU 上训练 TensorFlow 2 BERT 模型的最佳实践。

使用 BERT 的句子和句子对分类

句子和句子对分类任务指将给定的一对句子分类为“改述 (Paraphrases)”或“非改述 (Not paraphrases)”。以下是使用 Google Cloud Platform 上的多个 GPU 对 BERT-LARGE 模型进行微调的示例。此任务使用 Microsoft Research Paraphrase Corpus (MRPC) 语料库，其中包含 5,801 对句子以及人工释义，以指示每个句对是否存在改述/语义对等关系。

export BERT_BASE_DIR=gs://cloud-tpu-checkpoints/bert/keras_bert/uncased_L-24_H-1024_A-16
export MODEL_DIR=gs://some_bucket/my_output_dir
export GLUE_DIR=gs://some_bucket/datasets
export TASK=MRPC


python3 run_classifier.py \
  --mode='train_and_eval' \
  --input_meta_data_path=${GLUE_DIR}/${TASK}_meta_data \
  --train_data_path=${GLUE_DIR}/${TASK}_train.tf_record \
  --eval_data_path=${GLUE_DIR}/${TASK}_eval.tf_record \
  --bert_config_file=${BERT_BASE_DIR}/bert_config.json \
  --init_checkpoint=${BERT_BASE_DIR}/bert_model.ckpt \
  --train_batch_size=4 \
  --eval_batch_size=4 \
  --steps_per_loop=1 \
  --learning_rate=2e-5 \
  --num_train_epochs=3 \
  --model_dir=${MODEL_DIR} \
  --distribution_strategy=mirrored

与 ResNet 模型的 TPU 培训类似，用户可以通过使用 TPU 信息将分布策略类型更改为 tpu，从而轻松切换到 TPU 进行分布式训练，具体如以下示例所示：

--distribution_strategy=tpu
--tpu=grpc://${TPU_IP_ADDRESS}:8470

如您希望使用 TensorFlow Hub 中提供的预训练 SavedModel，只需将init_checkpointFLAG 替换为hub_module_urlFLAG，以指定 TensorFlow Hub 模块路径。

--hub_module_url=https://tfhub.dev/tensorflow/bert_en_uncased_L-24_H-1024_A-16/1

后续步骤

请访问 Model Garden 代码库，浏览我们在本文中介绍的示例代码。

在未来几个月里，我们将提供更多前沿的标准态模型和示例代码，以帮助您构建自己的模型。我们希望能有更多的 AI 研究人员和开发者在构建模型时使用 Model Garden。

我们也欢迎大家为代码库做出贡献，让整个 TensorFlow 社区都能从中受益。如您需要任何帮助，请在 GitHub 上联系我们。

如果您想详细了解 本文提及 的相关内容，请参阅以下文档。这些文档深入探讨了这篇文章中提及的许多主题：

• Model Garden (阅读原文) 
  https://github.com/tensorflow/models

• Model Garden 官方代码库
  https://github.com/tensorflow/models/tree/master/official

• TensorFlow Hub
  https://tfhub.dev/

• 分布策略 API 
  https://tensorflow.google.cn/guide/distribute_strategy

• MirroredStrategy 
  https://tensorflow.google.cn/api_docs/python/tf/distribute/MirroredStrategy
  

• MultiWorkerMirroredStrategy
  https://tensorflow.google.cn/api_docs/python/tf/distribute/experimental/MultiWorkerMirroredStrategy
  

• TPUStrategy

  https://tensorflow.google.cn/api_docs/python/tf/distribute/experimental/TPUStrategy

• ResNet 模型
  https://arxiv.org/abs/1512.03385

• 示例
  https://github.com/tensorflow/models/tree/master/official/vision/image_classification

•  tf.distribute.MirroredStrategy
  https://tensorflow.google.cn/api_docs/python/tf/distribute/MirroredStrategy

• classifier_trainer.py
  https://github.com/tensorflow/models/blob/master/official/vision/image_classification/classifier_trainer.py 

• compile
  https://tensorflow.google.cn/api_docs/python/tf/keras/Model#compile

• fit 
  https://tensorflow.google.cn/api_docs/python/tf/keras/Model#fit

• 示例
  https://github.com/tensorflow/models/blob/master/official/vision/image_classification/configs/examples/resnet/imagenet/gpu.yaml

• TF_CONFIG 
  https://tensorflow.google.cn/guide/distributed_training#TF_CONFIG

• TPU 配置文件
  https://github.com/tensorflow/models/blob/master/official/vision/image_classification/configs/examples/resnet/imagenet/tpu.yaml

• 自定义训练教程
  https://tensorflow.google.cn/tutorials/customization/custom_training

• BERT
  https://github.com/tensorflow/models/tree/master/official/nlp/bert

— 推荐阅读 —