使用 TensorFlow 2 进行训练

借助运行时版本 2.1 或更高版本在 AI Platform Training 上使用 TensorFlow 2 训练机器学习模型。TensorFlow 2 简化了许多来自 TensorFlow 1 的 API。TensorFlow 文档提供了将 TensorFlow 1 代码迁移到 TensorFlow 2 的指南。

在 AI Platform Training 中使用 TensorFlow 2 运行训练作业的过程与运行其他自定义代码训练作业的过程相同。但是,一些 AI Platform Training 功能在分别与 TensorFlow 2 和 TensorFlow 1 结合使用时,其工作方式存在差异。本文档对这些差异进行了总结。

Python 版本支持

运行时版本 2.1 及更高版本仅支持使用 Python 3.7 进行训练。因此,您必须借助 Python 3.7 来使用 TensorFlow 2 进行训练。

Python Software Foundation 于 2020 年 1 月 1 日终止了对 Python 2.7 的支持。2020 年 1 月 1 日之后发布的 AI Platform 运行时版本均不支持 Python 2.7。

分布式训练

TensorFlow 2 为分布式训练提供了更新的 API。此外,AI Platform Training 在运行时版本 2.1 及更高版本中设置不同的 TF_CONFIG 环境变量。本部分介绍了这两项变化。

分布策略

要在 TensorFlow 2 中针对多个虚拟机实例执行分布式训练,请使用 tf.distribute.Strategy API。我们尤其建议您将 Keras API 与 MultiWorkerMirroredStrategy 结合使用,或者,如果您为作业指定参数服务器,那么将 Keras API 与 ParameterServerStrategy 结合使用。不过请注意,TensorFlow 目前仅为这些策略提供实验性支持

TF_CONFIG

TensorFlow 要求在用于训练的每个虚拟机上设置 TF_CONFIG 环境变量。AI Platform Training 会自动在训练作业中使用的每个虚拟机上设置此环境变量。这样,每个虚拟机的行为都会因其类型的不同而有所不同,这有助于虚拟机相互通信。

在运行时版本 2.1 及更高版本中,AI Platform Training 不再在任何 TF_CONFIG 环境变量中使用 master 任务类型。而训练作业的主实例工作器TF_CONFIG 环境变量中标有 chief 类型。详细了解 AI Platform Training 如何设置 TF_CONFIG 环境变量

训练加速器

AI Platform Training 可让您使用 GPUTPU 来加速训练作业。

GPU

要了解如何使用 GPU 进行训练,请参阅 AI Platform Training 中的 GPU 配置指南TensorFlow 中的 GPU 使用指南

如果要在具有多个 GPU 的单个虚拟机上进行训练,最佳做法是使用 TensorFlow 的 MirroredStrategy。如果您要将多个虚拟机用于 GPU 进行训练,最佳做法是使用 TensorFlow 的 MultiWorkerMirroredStrategy

TPU

如需了解如何使用 TPU 进行训练,请参阅使用 TPU 进行训练指南

超参数调节

如果您使用 TensorFlow 2 运行超参数调节作业,则可能需要调整训练代码向 AI Platform Training 服务报告超参数调节指标的方式。

如果您是使用 Estimator 进行训练,您可以按 TensorFlow 1 中的方法将指标写入摘要。如果您是使用 Keras 进行训练,我们建议您使用 tf.summary.scalar 来撰写摘要。

详细了解如何报告超参数指标,并参阅在 TensorFlow 2 中如何执行此操作的示例。

后续步骤