发布者 Jaeyoun Kim，技术项目经理，和 Jing Li，软件工程师
我们想介绍一下 模型花园 的更新，该更新为 TensorFlow 用户提供了一个集中式位置，用于查找 TensorFlow 2 最先进模型的代码示例和可重用建模库。

模型花园旨在展示建模的最佳实践，以便 TensorFlow 用户可以充分利用最近的 TensorFlow 2 来进行研究和产品开发。TensorFlow Hub 将继续发挥其作为存储库的作用，用户可以在其中轻松搜索现成的预训练模型。我们计划越来越多地将模型花园中的更多最先进模型提供给 TensorFlow Hub，并且 TensorFlow Hub 页面将链接到模型花园中的模型实现。

我们最初提供了几个用于解决计算机视觉和自然语言处理问题的分布式训练的新代码示例。 模型花园官方存储库 中的代码示例将由 TensorFlow 团队维护，以使其与最新的 TensorFlow 2 API 保持同步。

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

模型花园中的许多模型可以以分布式方式进行训练。在 TensorFlow 2 中，您可以使用 分布策略 API 将训练工作负载分配到单主机/多加速器以及多主机/多加速器配置。以下是构建模型的一些常见分布策略，我们将在本博客中简要介绍。请访问 TensorFlow.org 以找到 TensorFlow 2 可用的分布策略列表。

• MirroredStrategy - 用于多个 GPU（如果未找到 GPU，则使用 CPU）
  
• MultiWorkerMirroredStrategy - 用于每个具有多个 GPU 的多个主机
  
• TPUStrategy - 用于多个 TPU 或多个 TPU 主机
  

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

计算机视觉

使用 ResNet 进行图像分类

ResNet 模型 可以将图像中的主要物体分类到 1,000 个物体类别中（例如，汽车、足球、台灯等）。模型花园提供了一个 示例，展示了如何在 GPU 和 TPU 上训练 ResNet。该模型是用纯粹的 TensorFlow 2 API 编写的，这些 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 数据集上这些值的示例可以在配置文件中找到（configs/examples/resnet/imagenet/gpu.yaml）。对于超参数调整，您可能想查看 Keras Tuner。

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

对于在多个 GPU 主机上的多个 GPU 上进行分布式训练，您需要在每个 GPU 主机上设置 TF_CONFIG 环境变量，以指定哪些任务构成集群、它们的地址以及每个任务在集群中的角色。以下是第一个 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” 部分指定了当前任务在集群中的角色。对于承担更多责任（例如，保存检查点）的主工作者，您需要将任务类型设置为“worker”，并将任务索引设置为 0。 MultiWorkerMirroredStrategy 将自动使用每个主机上所有可用的 GPU。

在云 TPU 上进行分布式训练

TensorFlow 2 使您能够轻松地在不同的硬件配置上进行分布式训练和训练模型，而无需更改模型定义。用户可以使用 tf.distribute.TPUStrategy 在云 TPU 上训练 ResNet 模型。您可以使用相同的框架（classifier_trainer.py）在云 TPU 上训练模型。您只需要使用不同的 TPU YAML 配置文件，并设置 --tpu=$TPU_NAME，其中 $TPU_NAME 是您在云控制台中 TPU 实例的名称。请参阅 TPU 配置文件（configs/examples/resnet/imagenet/tpu.yaml）。

$ 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 示例。模型花园包含 TensorFlow 2 对 BERT（来自 Transformers 的双向编码器表示）和 BERT 变体模型（例如，ALBERT）的实现。在这里，我们将演示在 GPU 和 TPU 上训练 TensorFlow 2 BERT 模型的最佳实践。

使用 BERT 进行句子和句子对分类

句子和句子对分类任务是将给定的句子对分类为释义或非释义。以下是在 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，请简单地用 hub_module_url FLAG 替换 init_checkpoint FLAG，以指定 TensorFlow Hub 模块路径。

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

后续步骤

请访问 模型花园存储库，查找本文中介绍的示例代码。

在接下来的几个月里，我们将提供更多最先进的规范模型和示例代码，用于构建您自己的模型。我们鼓励人工智能研究人员和开发人员在构建模型时利用模型花园。

我们也欢迎对存储库的贡献，以造福整个 TensorFlow 社区。如果您需要任何帮助，请在 GitHub 上与我们联系。