TensorFlow 用户指南

摘要

TensorFlow 是一个用于数值计算的开源软件库，它使用数据流图。图中的节点表示数学运算，而图的边表示在它们之间流动的多维数据数组（张量）。这种灵活的架构使您可以在桌面、服务器或移动设备中的一个或多个 CPU 或 GPU 上部署计算，而无需重写代码。《TensorFlow 用户指南》详细概述并深入介绍了如何使用和自定义 TensorFlow 深度学习框架。本指南还提供了有关 NVIDIA TensorFlow 参数的文档，您可以使用这些参数来帮助将容器的优化实现到您的环境中。

TensorFlow 是一个用于数值计算的开源软件库，它使用数据流图。

图中的节点表示数学运算，而图的边表示在它们之间流动的多维数据数组（张量）。这种灵活的架构使您可以在桌面、服务器或移动设备中的一个或多个 CPU 或 GPU 上部署计算，而无需重写代码。

TensorFlow 最初由 Google Brain 团队的研究人员和工程师在 Google 机器智能研究组织内开发，目的是进行机器学习和深度神经网络 (DNN) 研究。该系统具有足够的通用性，也适用于各种其他领域。

为了可视化 TensorFlow 结果，Docker^® 镜像还包含 TensorBoard。TensorBoard 是一套可视化工具。例如，您可以查看训练历史记录以及模型的样貌。

有关对 TensorFlow 所做的优化和更改的信息，请参阅《TensorFlow 发行说明》。

1.1. NVIDIA TensorFlow 容器的内容

此镜像包含 TensorFlow 的源代码和二进制文件。预构建和安装的 TensorFlow 版本位于 /usr/local/[bin,lib] 目录中。完整的源代码位于 /opt/tensorflow 中。

为了获得最佳的 TensorFlow 性能，容器镜像中包含示例脚本。有关更多信息，请参阅性能。

TensorFlow 包括 TensorBoard，这是一个由 Google 开发的数据可视化工具包。此外，此容器镜像还包含几个内置的 TensorFlow 示例，您可以使用如下命令运行这些示例。这些示例执行卷积神经网络 (CNN) 的训练。有关更多信息，请参阅 ML 初学者 MNIST 教程。以下 Python 命令运行其中两个示例

            

            
python -m tensorflow.models.image.mnist.convolutional
        

            

            
python -m tensorflow.models.image.cifar10.cifar10_multi_gpu_train

        

第一个命令使用 MNIST 数据集，例如，MNIST 数据库。第二个命令使用 CIFAR-10 数据集，例如，CIFAR-10 数据集。

要拉取 TensorFlow 容器，请参阅拉取容器

关于此任务

要运行 TensorFlow 容器，请参阅运行 TensorFlow。

验证 TensorFlow 是否正确运行的最简单方法是运行 /nvidia-examples/ 目录中包含的示例。

每个示例都包含一个 README 文件，其中描述了基本用法。

nvidia-docker 镜像经过预先打包、调整和准备运行；但是，您可能想要从头开始构建新镜像，或者使用自定义代码、库、数据或公司基础设施设置来扩充现有镜像。本节将指导您完成一些练习，这些练习将重点介绍如何从头开始创建容器、自定义容器、扩展深度学习框架以添加功能、从开发人员环境中使用该扩展框架开发一些代码，然后将该代码打包为版本化发布。

关于此任务

默认情况下，您不需要构建容器。NVIDIA 容器存储库 nvcr.io 有许多可以立即使用的容器，包括用于深度学习的容器以及仅包含 CUDA^® Toolkit™ 的容器。

容器的优点之一是它们可以用作创建新容器的起点。这可以称为自定义或扩展容器。您可以完全从头开始创建容器，但是，由于这些容器很可能在 GPU 上运行，因此建议您至少从包含操作系统和 CUDA^® 的 nvcr.io 容器开始。但是，您并不局限于此，并且可以创建在 CPU 上运行而不使用 GPU 的容器。在这种情况下，您可以从另一个位置（例如 Docker）的裸操作系统容器开始。为了简化开发，您仍然可以从带有 CUDA 的容器开始；只是在使用容器时不使用 CUDA。

自定义或扩展的容器可以保存到用户的私有容器存储库中。它们也可以与其他用户共享，但这需要一些管理员帮助。重要的是要注意，所有 NGC 深度学习框架镜像都包含构建框架本身以及所有先决条件的源代码。

最佳实践是避免使用 docker commit 来开发新的 Docker 镜像，而是使用 Dockerfile。 Dockerfile 方法提供了在 Docker 镜像开发期间所做更改的可视性和高效的版本控制能力。 Docker commit 方法仅适用于短期的、可丢弃的镜像。

有关编写 Docker 文件的更多信息，请参阅最佳实践文档。

5.1. 自定义 TensorFlow 的优势和局限性

您可以出于多种原因自定义容器以满足您的特定需求；例如，您依赖于 NVIDIA 提供的容器中未包含的特定软件。无论您的原因是什么，您都可以自定义容器。

除非容器镜像与框架源代码一起包含，否则它们不包含示例数据集或示例模型定义。请务必检查容器是否有示例数据集或模型。

5.2. 示例 1：使用 Dockerfile 自定义 TensorFlow

有关使用 Dockerfile 自定义 TensorFlow 的最新说明，请参阅容器内 /workspace/docker-examples/Dockerfile.customtensorflow 文件中提供的说明。
在自定义容器之前，您应确保已使用 docker pull 命令将 TensorFlow 21.08 容器加载到 NGC 容器注册表中，然后再继续。 例如

            

            
$ docker pull nvcr.io/nvidia/tensorflow:21.08
        

nvcr.io 上的 Docker 容器还提供了一个示例 Dockerfile，其中说明了如何修补框架并重建 Docker 镜像。在 /workspace/docker-examples 目录中，有两个您可以使用的示例 Dockerfile。第一个 Dockerfile.addpackages 可用于将软件包添加到 TensorFlow 镜像。第二个 Dockerfile.customtensorflow 说明了如何修补 TensorFlow 并重建镜像。

            

            
FROM nvcr.io/nvidia/tensorflow:21.08

# Bring in changes from outside container to /tmp
# (assumes my-tensorflow-modifications.patch is in same directory as Dockerfile)
COPY my-tensorflow-modifications.patch /tmp

# Change working directory to TensorFlow source path
WORKDIR /opt/tensorflow

# Apply modifications
RUN cd tensorflow-source \
  &amp;&amp; patch -p1 < /tmp/my-tensorflow-modifications.patch

# Rebuild TensorFlow
RUN ./nvbuild.sh

# Reset working directory
WORKDIR /workspace

        

此 DockerFile 将以与原始镜像中构建 TensorFlow 镜像相同的方式重建 TensorFlow 镜像。有关更多信息，请参阅Dockerfile 参考。为了更好地理解 Dockerfile，让我们逐步了解主要的命令。Dockerfile 中的第一行是以下内容

            

            
FROM nvcr.io/nvidia/tensorflow:21.08

        

此行以 NVIDIA 21.08 版本的 TensorFlow 镜像作为起点。
第二行是以下内容

            

            
COPY my-tensorflow-modifications.patch /tmp

        

它将容器外部的更改引入到您的 /tmp 目录中。这假定 my-tensorflow-modifications.patch 文件与 Dockerfile 位于同一目录中。
文件中接下来的重要行将工作目录更改为 TensorFlow 源代码路径。

            

            
WORKDIR /opt/tensorflow

        

接下来是用于将修改补丁应用于源代码的命令。

            

            
RUN cd tensorflow-source \
 &amp;&amp; patch -p1 < /tmp/my-tensorflow-modifications.patch

        

应用补丁后，可以重建 TensorFlow 镜像。这通过 DockerFile/ 中的 RUN 命令完成。

            

            
RUN ./nvbuild.sh

        

最后，DockerFile 中的最后一行主要行重置默认工作目录。

            

            
WORKDIR /workspace

        

5.3. 示例 2：使用 docker commit 自定义 TensorFlow

此示例使用 docker commit 命令将容器的当前状态刷新到 Docker 镜像。

关于此任务

这不是推荐的最佳实践，但是，当您有一个正在运行的容器，您在其中进行了更改并想要保存它们时，这很有用。在此示例中，我们使用 apt-get 标记来安装一个需要用户以 root 身份运行的软件包。

步骤

1. 从 nvcr.io 存储库将 Docker 容器拉取到您的 DGX™ 系统。例如，以下命令将拉取 TensorFlow 容器
   复制

   已复制！

               
   
               
   $ docker pull  nvcr.io/nvidia/tensorflow:21.07
           

2. 在您的 DGX 上运行容器。
   注意

   运行容器时，请勿使用 --rm 标志。如果在运行容器时使用 --rm 标志，则在退出容器时，您的更改将丢失。

   复制

   已复制！

               
   
               
   docker run --gpus all -ti nvcr.io/nvidia/tensorflow:21.07
           

   复制

   已复制！

               
   
               
   ================
   == TensorFlow ==
   ================
   
   NVIDIA Release 21.07 (build 31239)
   
   Container image Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
   Copyright 2021 The TensorFlow Authors.  All rights reserved.
   
   Various files include modifications (c) NVIDIA CORPORATION.  All rights reserved.
   NVIDIA modifications are covered by the license terms that apply to the underlying project or file.
   
   NOTE: The SHMEM allocation limit is set to the default of 64MB.  This may be
      insufficient for TensorFlow.  NVIDIA recommends the use of the following flags:
      docker run --gpus all --shm-size=1g --ulimit memlock=-1 --ulimit stack=67108864 ...
   
   root@8db6076d82c4:/workspace#
   
           

3. 您现在应该是容器中的 root 用户（注意提示符）。您可以使用 apt 命令来拉取软件包并将其放入容器中。
   注意

   NVIDIA 容器是使用 Ubuntu 构建的，Ubuntu 使用 apt-get 软件包管理器。在 深度学习文档中查看容器发行说明，以获取有关您正在使用的特定容器的详细信息。

   在此示例中，我们将 octave（MATLAB 的 GNU 克隆）安装到容器中。
   复制

   已复制！

               
   
               
   # apt-get update
   # apt install octave
           

   注意

   您必须先发出 apt-get update，然后再使用 apt 安装 Octave。

4. 退出工作区。
   复制

   已复制！

               
   
               
   # exit
           

5. 显示正在运行的容器列表。
   复制

   已复制！

               
   
               
   $ docker ps -a
           

   例如，以下是 docker ps -a 命令的一些输出
   复制

   已复制！

               
   
               
   $ docker ps -a
   CONTAINER ID    	IMAGE                         	       CREATED        ...      
   8db6076d82c4    	nvcr.io/nvidia/tensorflow:21.07     	3 minutes ago   ...	
   
           

6. 现在，您可以从正在运行的容器中创建一个新镜像，您已在其中安装了 Octave。您可以使用以下命令提交容器。
   复制

   已复制！

               
   
               
   $ docker commit 8db6076d82c4 nvcr.io/nvidian_sas/tensorflow_octave:21.07
   sha256:25198e37ae2e3416bebcf1d3084ff3a95600d978811fe7f4f184de0af3878b51
           

7. 显示镜像列表。
   复制

   已复制！

               
   
               
   $ docker images
   REPOSITORY                              TAG                           IMAGE ID       ...
   nvidian_sas/tensorflow_octave           21.07                         25198e37ae2e   ...
   
           

8. 为了验证，再次运行容器，看看 Octave 是否真的在那里。
   复制

   已复制！

               
   
               
   docker run --gpus all -ti nvidian_sas/tensorflow_octave:21.07
   
   ================
   == TensorFlow ==
   ================
   
   NVIDIA Release 21.07 (build 31239)
   
   Container image Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved. Copyright 2021 The TensorFlow Authors.  All rights reserved.
   
   Various files include modifications (c) NVIDIA CORPORATION.  All rights reserved. NVIDIA modifications are covered by the license terms that apply to the underlying project or file.
   
   NOTE: The SHMEM allocation limit is set to the default of 64MB.  This may be insufficient for TensorFlow.  NVIDIA recommends the use of the following flags:
      docker run --gpus all --shm-size=1g --ulimit memlock=-1 --ulimit stack=67108864 ...
   
   root@87e8dde4be6d:/workspace# octave
   octave: X11 DISPLAY environment variable not set
   octave: disabling GUI features
   GNU Octave, version 4.0.0
   Copyright (C) 2015 John W. Eaton and others.
   This is free software; see the source code for copying conditions.
   There is ABSOLUTELY NO WARRANTY; not even for MERCHANTABILITY or
   FITNESS FOR A PARTICULAR PURPOSE.  For details, type 'warranty'.
   
   Octave was configured for "x86_64-pc-linux-gnu".
   
   Additional information about Octave is available at http://www.octave.org.
   
   Please contribute if you find this software useful.
   For more information, visit http://www.octave.org/get-involved.html
   
   Read http://www.octave.org/bugs.html to learn how to submit bug reports.
   For information about changes from previous versions, type 'news'.
   
   octave:1>
   
           

   由于显示了 Octave 提示符，因此 Octave 已安装。

9. 如果您使用的是 DGX-1 或 DGX Station，并且想要将容器保存到您的私有存储库（Docker 使用短语“push”），则可以使用 docker push ... 命令。
   复制

   已复制！

               
   
               
   $ docker push nvcr.io/nvidian_sas/tensorflow_octave:21.07
           

5.4. 使用 TensorRT 加速 TensorFlow 中的推理

关于此任务

有关如何在 TensorFlow 框架中使用 TensorRT 的分步说明，请参阅《使用 TensorRT 加速 TensorFlow 推理用户指南》。要查看主要功能、软件增强和改进以及已知问题，请参阅《发行说明》。

NGC 容器注册表 (nvcr.io) 中的 TensorFlow 容器已按照以下参数的定义进行预配置。这些参数用于预编译 GPU，启用对加速线性代数 (XLA) 后端的支持，并禁用对 Google Cloud Platform (GCP) 和 Hadoop 分布式文件系统 (HDFS) 的支持。

6.1. 添加和修改的参数

除了 Google TensorFlow 容器中包含的 Dockerfile 中的参数外，以下参数已在 NVIDIA 版本的 TensorFlow 中添加或修改。

对于本指南中未提及的参数，请参阅TensorFlow 文档。

6.1.1. TF_CUDA_COMPUTE_CAPABILITIES

TF_CUDA_COMPUTE_CAPABILITIES 参数使代码能够针对特定的 GPU 架构进行预编译。

该容器构建时带有以下设置，该设置针对 Pascal、Volta、Turing 和 NVIDIA Ampere 架构 GPU

            

            
TF_CUDA_COMPUTE_CAPABILITIES "6.0,6.1,7.0,7.5,8.0,8.6"

        

其中数字对应于 GPU 架构

6.1.2. TF_NEED_GCP

TF_NEED_GCP 参数按定义禁用对 Google Cloud Platform (GCP) 的支持。

该容器构建时带有以下设置，该设置关闭对 GCP 的支持

            

            
TF_NEED_GCP 0

        

6.1.3. TF_NEED_HDFS

TF_NEED_HDFS 参数按定义禁用对 Hadoop 分布式文件系统 (HDFS) 的支持。

该容器构建时带有以下设置，该设置关闭对 HDFS 的支持

            

            
TF_NEED_HDFS 0

        

6.1.4. TF_ENABLE_XLA

TF_ENABLE_XLA 参数按定义启用对加速线性代数 (XLA) 后端的支持。

该容器构建时带有以下设置，该设置开启对 XLA 的支持

            

            
TF_ENABLE_XLA 1

        

以下环境变量设置在 TensorFlow 中启用某些功能。它们会稍微改变和降低计算的精度，并且默认情况下处于启用状态。

7.1. 添加或修改的变量

除了 Google TensorFlow 容器中包含的 Dockerfile 中的变量外，以下变量已在 NVIDIA 版本的 TensorFlow 中添加或修改。

对于本指南中未提及的变量，请参阅TensorFlow 文档。

7.1.1. TF_ADJUST_HUE_FUSED

TF_ADJUST_HUE_FUSED 变量启用融合内核以用于图像色调。
此变量默认启用

            

            
TF_ADJUST_HUE_FUSED         1

        

要禁用该变量，请运行以下命令

            

            
export TF_ADJUST_HUE_FUSED=0

        

7.1.2. TF_ADJUST_SATURATION_FUSED

TF_ADJUST_SATURATION_FUSED 变量启用融合内核以用于饱和度调整。
此变量默认启用

            

            
TF_ADJUST_SATURATION_FUSED  1

        

要禁用该变量，请运行以下命令

            

            
export TF_ADJUST_SATURATION_FUSED=0

        

7.1.3. TF_ENABLE_WINOGRAD_NONFUSED

TF_ENABLE_WINOGRAD_NONFUSED 变量启用非融合 Winograd 卷积算法的使用。
此变量默认启用

            

            
TF_ENABLE_WINOGRAD_NONFUSED 1

        

要禁用该变量，请运行以下命令

            

            
export TF_ENABLE_WINOGRAD_NONFUSED=0

        

7.1.4. TF_AUTOTUNE_THRESHOLD

TF_AUTOTUNE_THRESHOLD 变量提高了用于选择最快卷积算法的自动调整过程的稳定性。将其设置为更高的值可以提高稳定性，但需要在训练开始时进行更多试验步骤，然后才能找到最佳算法。
在容器中，此变量设置为以下值

            

            
export TF_AUTOTUNE_THRESHOLD=2

        

要将此变量设置为其默认设置，请运行以下命令

            

            
export TF_AUTOTUNE_THRESHOLD=1

        

7.1.5. CUDA_DEVICE_MAX_CONNECTIONS

CUDA_DEVICE_MAX_CONNECTIONS 变量解决了与 Tesla K80 GPU 上的流相关联的性能问题。
在容器中，此变量设置为以下值

            

            
export CUDA_DEVICE_MAX_CONNECTIONS=12

        

要将此变量设置为其默认设置，请运行以下命令

            

            
export CUDA_DEVICE_MAX_CONNECTIONS=8

        

7.1.6. TF_DISABLE_CUDNN_TENSOR_OP_MATH

TF_DISABLE_CUDNN_TENSOR_OP_MATH 变量启用和禁用 Tensor Core 数学运算，以用于 TensorFlow 中的 cuDNN 卷积。

默认情况下启用 Tensor Core 数学运算，但可以通过将此变量设置为 1 来禁用。有关更多信息，请参阅Tensor Core 数学运算。默认情况下禁用此变量

            

            
export TF_DISABLE_CUDNN_TENSOR_OP_MATH=0

        

要启用该变量，请运行以下命令

            

            
export TF_DISABLE_CUDNN_TENSOR_OP_MATH=1

        

7.1.7. TF_DISABLE_CUDNN_RNN_TENSOR_OP_MATH

TF_DISABLE_CUDNN_RNN_TENSOR_OP_MATH 变量启用和禁用 Tensor Core 数学运算，以用于 TensorFlow 中的 cuDNN RNN。

默认情况下启用 Tensor Core 数学运算，但可以通过将此变量设置为 1 来禁用。有关更多信息，请参阅Tensor Core 数学运算。默认情况下禁用此变量

            

            
export TF_DISABLE_CUDNN_RNN_TENSOR_OP_MATH=0
        

要启用该变量，请运行以下命令

            

            
export TF_DISABLE_CUDNN_RNN_TENSOR_OP_MATH=1
        

7.1.8. TF_DISABLE_CUBLAS_TENSOR_OP_MATH

TF_DISABLE_CUBLAS_TENSOR_OP_MATH 变量启用和禁用 Tensor Core 数学运算，以用于 TensorFlow 中的 cuBLAS 卷积。

默认情况下启用 Tensor Core 数学运算，但可以通过将此变量设置为 1 来禁用。有关更多信息，请参阅Tensor Core 数学运算。默认情况下禁用此变量

            

            
export TF_DISABLE_CUBLAS_TENSOR_OP_MATH=0

        

要启用该变量，请运行以下命令

            

            
export TF_DISABLE_CUBLAS_TENSOR_OP_MATH=1

        

7.1.9. TF_ENABLE_CUBLAS_TENSOR_OP_MATH_FP32

TF_ENABLE_CUBLAS_TENSOR_OP_MATH_FP32 变量启用和禁用 Tensor Core 数学运算，以用于 TensorFlow 中的 float32 矩阵乘法运算。

默认情况下禁用用于 float32 运算的 Tensor Core 数学运算，但可以通过将此变量设置为 1 来启用。有关更多信息，请参阅Tensor Core 数学运算。默认情况下禁用此变量

            

            
export TF_ENABLE_CUBLAS_TENSOR_OP_MATH_FP32=0

        

要启用此变量，请运行以下命令

            

            
export TF_ENABLE_CUBLAS_TENSOR_OP_MATH_FP32=1

        

7.1.10. TF_ENABLE_CUDNN_TENSOR_OP_MATH_FP32

TF_ENABLE_CUDNN_TENSOR_OP_MATH_FP32 变量启用和禁用 Tensor Core 数学运算，以用于 TensorFlow 中的 float32 卷积运算。

默认情况下禁用用于 float32 运算的 Tensor Core 数学运算，但可以通过将此变量设置为 1 来启用。有关更多信息，请参阅Tensor Core 数学运算。默认情况下禁用此变量

            

            
export TF_ENABLE_CUDNN_TENSOR_OP_MATH_FP32=0

        

要启用此变量，请运行以下命令

            

            
export TF_ENABLE_CUDNN_TENSOR_OP_MATH_FP32=1

        

7.1.11. TF_ENABLE_CUDNN_RNN_TENSOR_OP_MATH_FP32

TF_ENABLE_CUDNN_RNN_TENSOR_OP_MATH_FP32 变量启用和禁用 Tensor Core 数学运算，以用于 TensorFlow 中的 float32 cuDNN RNN 运算。默认情况下禁用用于 float32 运算的 Tensor Core 数学运算，但可以通过将此变量设置为 1 来启用。有关更多信息，请参阅Tensor Core 数学运算。
默认情况下禁用此变量

            

            
export TF_ENABLE_CUDNN_RNN_TENSOR_OP_MATH_FP32=0

        

要启用此变量，请运行以下命令

            

            
export TF_ENABLE_CUDNN_RNN_TENSOR_OP_MATH_FP32=1

        

7.1.12. TF_ENABLE_LAYOUT_NHWC

TF_ENABLE_LAYOUT_NHWC 变量在模型执行期间强制执行 NHWC (channels_last) 格式。它主要适用于 float32 和 NVIDIA Ampere 架构 GPU 或更高版本，在这些版本中将使用 TF32 Tensor Core，并且首选 NHWC 格式。

默认情况下禁用此环境变量。禁用后，数据格式将由 TensorFlow 自动选择。要启用该变量，请运行以下命令

            

            
export TF_ENABLE_LAYOUT_NHWC=1
        

7.1.13. TF_ENABLE_NVTX_RANGES

NVIDIA Tools Extension (NVTX) 范围在通过 Nsight Systems 或 NVIDIA Visual Profiler 分析应用程序时，将操作名称注释添加到执行时间线。

有关 NVTX 的更多信息，请参阅https://docs.nvda.net.cn/cuda/profiler-users-guide/index.html#nvtx

自 21.03 版本以来的说明

TF_ENABLE_NVTX_RANGES 变量启用和禁用 TensorFlow 中的 NVTX 范围。默认情况下禁用 NVTX 范围，但可以通过将此变量设置为 1 来启用。默认情况下禁用此变量

            

            
export TF_ENABLE_NVTX_RANGES=0
        

要启用该变量，请运行以下命令

            

            
export TF_ENABLE_NVTX_RANGES=1
        

除了设置 TF_ENABLE_NVTX_RANGES=1 之外，还可以将变量 TF_ENABLE_NVTX_RANGES_DETAILED 设置为 1，以获取包含更详细信息（例如形状和属性）的 NVTX 范围。默认情况下禁用此变量

            

            
export TF_ENABLE_NVTX_RANGES_DETAILED=0
        

要启用该变量，请运行以下命令

            

            
export TF_ENABLE_NVTX_RANGES_DETAILED=1
        

请注意，不再支持 TF_DISABLE_NVTX_RANGES 变量（见下文）。

21.03 版本之前的说明 TF_ENABLE_NVTX_RANGES 变量默认禁用

            

            
export TF_DISABLE_NVTX_RANGES=0
        

要启用此变量，请运行以下命令

            

            
export TF_DISABLE_NVTX_RANGES=1
        

7.1.14. TF_CUDNN_CTC_LOSS

TF_CUDNN_CTC_LOSS 变量通过 nn.ctc_loss 为 Tensorflow 2.x（例如，19.11-tf2-py3）或 nn.ctc_loss_v2 为 Tensorflow 1.x（例如，19.11-tf1-py3）启用 cuDNN CTC 损失后端。
默认情况下禁用此变量

            

            
export TF_CUDNN_CTC_LOSS=0
        

要启用此变量，请运行以下命令

            

            
export TF_CUDNN_CTC_LOSS=1
        

7.1.15. TF_GPU_ALLOCATOR

TF_GPU_ALLOCATOR 变量使用自 CUDA 11.2 起可用的 cudaMallocAsync 启用内存分配器。与默认的 BFC 内存分配器相比，它的碎片问题更少。请注意，这将来可能会成为默认设置。

默认情况下未设置此环境变量。要启用该变量，请运行以下命令

            

            
export TF_GPU_ALLOCATOR=cuda_malloc_async
        

7.1.16. TF_GRAPPLER_GRAPH_DEF_PATH

TF_GRAPPLER_GRAPH_DEF_PATH 变量输出 TF grappler 优化之前和之后的 Graphdef 文件（有关 grappler 优化的更多信息，请参阅使用 Grappler 进行 TensorFlow 图优化）。在检查 TF 运行时期间优化的操作图时，用户可以指定 TF_GRAPPLER_GRAPH_DEF_PATH=/<i>path</i>/<i>to</i>/<i>graphdef</i>。

默认情况下未设置此环境变量。要启用该变量，请运行以下命令

            

            
export TF_GRAPPLER_GRAPH_DEF_PATH="path/to/graphdef"
        

为了获得最佳的 TensorFlow 性能，对于基于图像的训练，容器包含示例脚本，演示了 CNN 的高效训练。示例脚本可能需要修改以适应您的应用程序。

脚本可以在 /opt/tensorflow/nvidia-examples/cnn/ 目录中找到。除了训练脚本外，还可以在 /opt/tensorflow/nvidia-examples/cnn/README.md 目录中找到一些文档。

有关更多信息，请参阅性能模型和基准测试。

8.1. Tensor Core 数学运算

TensorFlow 容器包含对 Volta 架构（在 Tesla V100 GPU 上可用）中 Tensor Core 的支持。Tensor Core 为训练提供高达 12 倍的峰值 TFLOP。容器默认启用 Tensor Core 数学运算；因此，任何包含使用 tf.float16 数据类型的卷积或矩阵乘法的模型都将在可能的情况下自动利用 Tensor Core 硬件。

对于 tf.float32 矩阵乘法、卷积和 RNN 运算，也可以通过分别设置 TF_ENABLE_CUBLAS_TENSOR_OP_MATH_FP32=1, TF_ENABLE_CUDNN_TENSOR_OP_MATH_FP32=1 和 TF_ENABLE_CUDNN_RNN_TENSOR_OP_MATH_FP32=1 （对于使用 cudnn_rnn op 的 RNN）环境变量来启用 Tensor Core 数学运算。此模式会导致数据在内部降至 float16 精度，这可能会影响训练收敛。

启用 Tensor Core 数学运算后，矩阵乘法、卷积和 RNN 运算的输入会从 FP32 隐式向下转换为 FP16。内部累积和输出仍为 FP32。这允许 FP32 模型通过在可用时使用 GPU Tensor Core 来更快地运行。此外，用户应扩充模型以包含损失缩放（例如，通过将优化器包装在 tf.contrib.mixed_precision.loss_scale_optimizer 中）。

有关架构的更多信息，请参阅Volta 内部和Turing 内部。

8.1.1. Float16 训练

在某些情况下，使用降低的精度进行训练可能会导致收敛不良或不稳定。
NVIDIA 建议使用以下策略来最大限度地减少训练期间降低精度的影响（有关 float16 训练的完整演示，请参阅 nvidia-examples/cnn/nvcnn.py）

1. 将可训练变量保持在 float32 精度，并在模型中使用它们之前将其转换为 float16。例如
   复制

   已复制！

               
   
               
   tf.cast(tf.get_variable(..., dtype=tf.float32), tf.float16)
           

2. 如果模型难以收敛或无法收敛，请应用损失缩放。损失缩放涉及在计算梯度之前将损失乘以比例因子，然后再次将所得梯度除以相同的比例以重新标准化它们。循环神经网络模型的典型损失比例因子为 128。例如
   复制

   已复制！

               
   
               
   loss, params = ...
   scale = 128
   grads = [grad / scale for grad in tf.gradients(loss * scale, params)]
           

8.2. 自动混合精度 (AMP)

使用混合精度训练需要三个步骤

1. 将模型转换为尽可能使用 float16 数据类型。
2. 保留 float32 主权重以累积每次迭代的权重更新。
3. 使用损失缩放来保留小的梯度值。

在 TensorFlow 框架中使用自动混合精度可以像添加一行代码一样简单。它通过自动重写所有计算图以及必要的运算来实现混合精度训练和损失缩放。有关更多信息，请参阅深度学习的自动混合精度。

8.2.1. TensorFlow 中的自动混合精度训练

对于已经使用 tf.train 或 tf.keras.optimizers 中的优化器进行 compute_gradients() 和 apply_gradients() 操作的模型（例如，通过调用 optimizer.minimize() 或 model.fit()），可以通过使用 tf.train.experimental.enable_mixed_precision_graph_rewrite() 包装优化器来启用自动混合精度。
基于图的示例

            

            
opt = tf.train.AdamOptimizer()
opt = tf.train.experimental.enable_mixed_precision_graph_rewrite(opt)
train_op = opt.miminize(loss)
        

基于 Keras 的示例

            

            
opt = tf.keras.optimizers.Adam()
opt = tf.train.experimental.enable_mixed_precision_graph_rewrite(opt)
model.compile(loss=loss, optimizer=opt)
model.fit(...)
        

有关此功能的更多信息，请参阅 TensorFlow 文档此处。为了向后兼容以前的容器版本，还可以通过定义以下环境变量为 tf.train 优化器启用 AMP

            

            
export TF_ENABLE_AUTO_MIXED_PRECISION=1
        

启用后，自动混合精度将执行以下两项操作：

1. 在您的 TensorFlow 图中插入适当的类型转换操作，以在适当的位置使用 float16 执行和存储 - 这可以使用 Tensor Core 以及节省内存存储和带宽。
2. 打开训练优化器对象内部的自动损失缩放。

8.2.2. 条件和限制

确保您熟悉以下条件

其他控制

可以在不进行自动损失缩放的情况下启用自动插入类型转换操作。用于执行此操作的环境变量是 TF_ENABLE_AUTO_MIXED_PRECISION_GRAPH_REWRITE=1。

注意事项

模型类型

主要依赖于分组卷积或深度可分离卷积的卷积架构（MobileNet 和 ResNeXt 是流行的示例）目前不会从 float16 执行中看到加速。这是由于超出自动混合精度范围的库约束，但我们预计它们很快会放宽。

优化器

自动混合精度损失缩放要求模型代码使用内置 tf.train.optimizer 或 tf.keras.optimizers.optimizer 类的子类。此外

• 直接调用 tf.gradients 并“手动”使用这些梯度的 TensorFlow 代码将不受支持。
• 相反，自动混合精度需要成对调用 optimizer.compute_gradients 和 optimizer.apply_gradients，或者调用高级函数 optimizer.minimize。
• 如果优化器类是 tf.train.optimizer 的自定义子类（不是 TensorFlow 内置的子类），则自动混合精度损失缩放可能不支持它。特别是，如果自定义子类覆盖 compute_gradients 或 apply_gradients 中的任何一个，则必须注意也要调用这些方法的超类实现。

多 GPU

在 TensorFlow 1.14.0 之前，自动混合精度不支持 TensorFlow “分布式策略”。相反，多 GPU 训练需要使用 Horovod（或 TensorFlow 设备原语）。

其他说明

如果您的代码已经内置了自动损失缩放支持，则需要禁用它，以避免与自动混合精度自身的自动损失缩放冲突。或者，可以在不启用损失缩放的情况下启用自动混合精度图重写，方法是使用上面描述的选项。

8.2.3. 常见问题解答

问：如果我的模型代码已经支持混合精度训练怎么办？

如果代码已经编写为遵循《混合精度训练指南》，则自动混合精度将保持原样。例如，NVIDIA TensorFlow 容器内提供的 CNN 示例默认使用混合精度训练。如果您想评估它们在自动混合精度下的工作方式，请务必使用标志 --precision=fp32 运行它们。

问：使用自动混合精度后，我的模型会快多少？

混合精度加速没有精确的规则，但以下是一些指导原则

• 在矩阵乘法（密集层）或卷积中花费的时间越多，Tensor Core 可以加速模型的程度就越高。这意味着“更大”的模型通常会看到更大的加速。特别是，非常小的密集层和卷积层将从自动混合精度中看到有限的好处，因为没有足够的数学运算来充分利用 Tensor Core。
• 混合精度模型比 FP32 使用更少的内存，因此在使用自动混合精度运行时可以增加批量大小。因此，您通常可以通过在启用自动混合精度后增加批量大小来提高加速。

问：如何查看自动混合精度对我的模型所做的更改？

由于自动混合精度在 TensorFlow 图的级别上运行，因此快速掌握它所做的更改可能具有挑战性：它通常会调整数千个 TensorFlow 操作，但这些操作对应于少得多的逻辑层。您可以设置环境变量 TF_CPP_VMODULE="auto_mixed_precision=2" 以查看自动混合精度所做决定的完整日志（请注意，这可能会生成大量输出）。

问：为什么我在保存的模型 GraphDef 中只看到 FP32 数据类型？

当您保存模型图或使用 Session.graph 检查 Session.graph_def 的图时，TensorFlow 会返回未优化的图版本。自动混合精度充当原始图上的优化过程，因此其更改未包含在未优化的图中。您可以设置环境变量 TF_AUTO_MIXED_PRECISION_GRAPH_REWRITE_LOG_PATH=”my/log/path”，自动混合精度会将它处理的每个图的优化前和优化后副本保存到该目录。

问：为什么在使用自动混合精度进行训练时，我看到 step=0 重复多次？

自动混合精度启用的自动损失缩放算法可以选择在搜索最佳损失比例时“跳过”训练迭代。当它这样做时，它不会增加全局步数计数。由于大多数跳过发生在训练开始时（通常少于十次迭代），因此此行为表现为步数计数器保持为零的多次迭代。

问：如何处理用户定义的自定义 TensorFlow 操作？

默认情况下，自动混合精度将保留任何它不了解的操作类型，包括自定义操作。这意味着操作的输入和输出类型不会更改，并且自动混合精度将根据需要插入类型转换，以便与（可能已更改的）图的其余部分互操作。如果您希望自动混合精度识别自定义操作类型，可以使用三个环境变量

TF_AUTO_MIXED_PRECISION_GRAPH_REWRITE_ALLOWLIST_ADD

这些操作值得将输入转换为 FP16 以获得 FP16 执行。 它们主要是可以利用 Tensor Core 的操作。

TF_AUTO_MIXED_PRECISION_GRAPH_REWRITE_INFERLIST_ADD

这些操作可以使用 FP16 执行，因此如果输入恰好已经是 FP16（由于上游 ALLOWLIST 操作），则它们可以使用 FP16。

TF_AUTO_MIXED_PRECISION_GRAPH_REWRITE_DENYLIST_ADD

这些操作需要 FP32 才能保证数值精度，并且输出不能安全地转换回 FP16。 示例操作包括 Exp 和 Log。

这些环境变量中的每一个都接受逗号分隔的字符串操作名称列表。 例如，您可以设置 export TF_AUTO_MIXED_PRECISION_GRAPH_REWRITE_ALLOWLIST_ADD=MyOp1,MyOp2。 操作名称是在调用 REGISTER_OP 时使用的字符串名称，它对应于操作 OpDef 上的 name 属性。

问：我可以更改自动混合精度的算法行为吗？

控制自动混合精度行为的主要手段是操作允许列表、推断列表和拒绝列表中的操作。 您可以使用上述三个环境变量将操作添加到每个列表，并且有一个对应的变量 TF_AUTO_MIXED_PRECISION_GRAPH_REWRITE_{ALLOWLIST,INFERLIST,DENYLIST}_REMOVE 可以从每个列表中删除内置操作。

9.1. XLA 简介

XLA 是 TensorFlow 的优化图编译器。 它优化 TensorFlow GraphDef 的某些部分，以尝试提高性能。

与一次执行一个 GraphDef 节点的原生 TensorFlow 不同，XLA 一次考虑多个 GraphDef 节点，并为这些节点生成优化的代码。

9.1.1. 为什么使用 XLA？

在许多情况下，XLA 比原生 TensorFlow 提高了性能。 两者之间的主要区别在于 XLA 中的融合优化器。 XLA 将许多小型内核优化为更大的内核，而不是背靠背地执行许多小型内核。 这大大减少了带宽受限内核的执行时间。 XLA 还提供了许多代数简化，远远优于 Tensorflow 提供的简化。

9.1.2. 启用 XLA

可以通过几种方式启用 XLA

• 选择加入模型的特定部分，称为手动集群。

  tf.function 是用户控制使用 XLA 优化模型哪些部分的一种方式。 有关 TF1.X，请参阅此处，有关 TF2.X，请参阅此处，了解如何使用 tf.function。

  或者，可以使用所谓的 jit_scope 来控制使用 XLA 优化模型的哪些部分。 有关 TF1.X，请参阅此处，有关 TF2.X，请参阅此处，了解如何使用 jit_scope。 请注意，这仅仅是一个提示，因为不支持的节点不会被编译。

  让 TensorFlow 决定优化模型的哪些部分，称为自动集群。

• 在源代码级别选择加入。

  对于 TF 1.X，用户可以通过在会话配置参数中添加以下代码行，在会话级别选择加入自动集群

  config.graph_options.optimizer_options.global_jit_level =

  tf.OptimizerOptions.ON_1

  对于 TF 2.X，用户可以通过在脚本开头添加以下代码行，在全局级别选择加入自动集群

  tf.config.optimizer.set_jit(True)

• 通过环境变量选择加入。

  用户可以无需任何源代码更改即可选择加入自动集群，其中 TensorFlow 决定哪些节点进入 XLA，方法是设置以下环境变量

  TF_XLA_FLAGS=--tf_xla_auto_jit=1

  此环境变量适用于 TF 1.X 和 TF 2.X。 对于通过会话配置启用 XLA 的模型，可以使用以下命令禁用 XLA

  TF_XLA_FLAGS=--tf_xla_auto_jit=-1

环境变量优先于 global_jit_level 和 jit_scope，但不优先于 tf.function。 本文档主要关注自动集群。 大部分内容也适用于手动集群。

9.1.2.1. XLA Lite

TensorFlow 自动集群根据允许列表对操作进行集群。

默认情况下，此列表非常广泛。 这可能会导致大型集群，从而导致性能下降或过度内存使用（见下文）。 为了缓解此类问题，用户可以设置以下内容，使 TensorFlow 仅对“小型”和可融合的操作（如逐点和归约操作）进行集群

TF_XLA_FLAGS=--tf_xla_auto_jit=fusible

请注意，由于模型始终具有矩阵乘法或卷积，因此这会导致更多但更小的集群。

9.1.3. XLA 注意事项

与原生 TensorFlow 执行相比，使用 XLA 会产生两种不同的成本

1. 编译时间

   代码必须在运行时编译。 这需要时间，具体取决于生成的集群的大小以及编译的次数（每个形状实例一次），此时间可能无法在执行期间恢复。

2. 由于 TF-XLA 接口导致的执行时间开销。

   XLA 集成回 TensorFlow 的方式增加了图执行的成本。

本文档通篇更深入地讨论了这两点。 此处提到这些注意事项是为了强调 XLA 并非适用于所有场景下加速所有模型的灵丹妙药。 短时间运行的脚本（小批量大小）通常不是 XLA 的良好候选对象。 表现出动态形状的模型是另一类在使用 XLA 时通常表现不佳的模型。

此外，使用 XLA 产生的开销很难确定。 有些开销发生在模型开始时，有些开销在整个执行过程中以固定成本发生，还有一些开销在执行过程中不规则地发生。 在指示代码从模型执行中获取计时信息时，务必注意这一点。

除了这些性能注意事项外，XLA 有时还会导致执行期间内存不足。 这将在后面更详细地讨论。

9.2. TF-XLA 集成

本节介绍 XLA 当前如何更改执行的图，以及它如何影响图的执行。

这有助于开发人员了解 TensorFlow 如何与 XLA 接口，以及这如何影响性能。 本节以及本文档的其余部分假定读者对 TensorFlow GraphDef 以及 TensorFlow 流执行器如何执行它有基本的了解。

9.2.1. TensorFlow 图的更改

9.2.1.1. 集群

启用自动集群后，将选择图的一部分与 XLA 一起编译。 这部分 (G) 具有一组输入 (I) 和一组输出 (O)。

启用 XLA 后，此图将转换为以下图以供执行

部分 (G) 已被集群到集群 (C) 中。 请注意，C 仅仅是操作的集群，尚未编译。 它还不能编译，因为输入和输出张量的实际形状在执行时才可用。

当自动集群时，模型的每个集群部分都将替换为上面描述的图，该图引用表示该部分的集群。 请注意，图的原始 GraphDef 表示仍然存在，以用作回退路径。 合并节点（在实际图中，每个输出有一个合并节点）从原生 TensorFlow 执行或 XLA 执行转发输出。

9.2.1.2. 使用 XLA 的 TensorFlow 图执行

在图执行期间，当到达上述图时，对于每组唯一的输入形状实例，都会发生以下情况。¹

1. 首次执行 _XlaCompile 节点的前两次，它会回退到 G 的原始 GraphDef 执行。
2. 第三次执行 _XlaCompile 节点时，集群 C 将编译为 XLA 二进制文件。 编译完成后，二进制文件将存储在缓存中，可通过对集群进行编码的密钥以及输入的形状和类型进行访问。 此密钥传递到 _XlaRun 节点，并执行二进制文件。 如果集群的编译时间超过 30 秒，则将来不会为任何其他形状实例编译它。
3. 后续每次执行 _XlaCompile 节点时，都会跳过编译，并将缓存二进制文件的密钥传递到 _XlaRun 节点以执行。

   请注意，为了使 _XlaCompile 节点执行，所有输入 (I) 都必须准备就绪。 同样，在 G 或 _XlaRun 完成执行之前，任何输出 (O) 都不会准备就绪。

   请注意，以上步骤适用于每个唯一的形状实例。 如果任何输入的形状在模型执行过程中发生变化，则会发生集群的重新编译。

9.2.2. XLA 问题的症状

本节提到使用 XLA 时观察到最常见的问题，以及这些问题发生的原因。 下一节重点介绍如何解决这些问题。

9.2.2.1. 功能问题

最常见的功能问题分为两类

1. 不同的输出。

   XLA 是一种编译器，仅执行语义等效的转换。² 由于这些转换是在浮点张量上进行的，因此与原生 TensorFlow 执行相比，结果很可能（略有）不同。 开发人员应该意识到这一点。

2. 内存不足。

   XLA 由于其设计本质，增加了执行所需的内存量。 如前所述，所有输入都必须在执行前准备就绪，并且所有输出都在执行后同时准备就绪。 这意味着必须为执行分配所有输入和输出的内存。 此外，还必须分配 XLA 二进制文件中保存中间张量所需的所有内存。

9.2.2.2. 性能问题

最常见的性能问题分为四类

1. TF-XLA 集成。

   TF 与 XLA 接口的方式（如 9.2.1 中所述）会带来性能成本，部分原因是图中插入了额外的节点。

2. 编译时间开销。

   当模型具有不同形状的输入数据时（例如 NLP 模型的句子），编译开销会累积。 大型集群也对此有贡献。

3. 计算/通信重叠。

   XLA 要求所有输入在执行前准备就绪，并且输出在执行完成时准备就绪。 与 Tensorflow 相比，这种同步可能会导致性能下降。 TensorFlow 具有细粒度的执行模型，允许复制和计算重叠，这在使用 XLA 时是不可能的。

4. XLA 优化和代码生成。

   有时 XLA 优化器或代码生成器不如他们本可以做到的那样好。 根据我们的经验，这是处理性能回归时最不常见的情况。

9.3. 识别和管理问题

对于开发人员来说，XLA 大多是一个黑匣子。 本节阐明了如何识别症状，并采取适当的措施来尝试解决这些问题。

9.3.1. 使用环境变量控制 XLA

XLA 公开了许多方法来从中检索信息或控制其行为。 这些选项没有（很好地）记录在案，甚至可能在版本之间发生变化（行为、格式甚至存在）。 即使这样，它们也提供了对 XLA 内部结构的大量见解。 它们分为两类

• TF_XLA_FLAGS 环境变量

  此变量可以设置为控制 TF-XLA 边界的选项。 它可以设置为以空格分隔的选项列表，可以在 tensorflow/compiler/jit/flags.cc 中找到

• XLA_FLAGS 环境变量

  此变量可以设置为控制 XLA 优化器和代码生成器的选项。 它可以设置为以空格分隔的选项列表，可以在 tensorflow/compiler/xla/debug_options_flags.cc 中找到

这些源文件中的大多数选项主要供 XLA 开发人员使用，并且超出了本文档的范围。 其中一些对模型开发人员很有用，并在下面提到。

9.3.2. 内存不足问题

这很容易识别，因为它是由 TensorFlow 分配器报告的。 内存不足的原因已在功能问题中描述。 发生这种情况时，请使用此额外参数运行模型

TF_XLA_FLAGS=--tf_xla_always_defer_compilation=true³

此选项指示 TensorFlow 永远不要编译集群，而是始终执行回退路径。 如果这仍然导致内存不足错误，则问题在于输入 (I) 和输出 (O) 无法同时全部放入内存。 如果执行成功，则问题在于输入 (I)、输出 (O) 和 XLA 二进制文件中的所有中间张量无法同时全部放入内存。

在这两种情况下，保证解决 XLA 中内存不足问题的唯一方法是减少集群中的操作数量。 操作数量（间接）影响输入、输出和中间张量的数量。

减少集群大小的一种方法是启用 XLA-Lite。 这几乎保证了此问题会消失，因为集群大小已大大减小。

另一种方法是限制每个集群的最大大小。 默认情况下，集群的大小没有上限。 可以使用以下命令检索集群的大小

TF_CPP_VMODULE=mark_for_compilation_pass=2

查找 " *** Clustering info for graph" 的出现次数，以获取生成的集群的大小。

当使用以下命令运行时

TF_XLA_FLAGS=--tf_xla_max_cluster_size=<n>

其中 <n> 小于最大集群大小，可以找到避免内存不足的上限。

减少小批量大小是另一种可以避免内存不足问题的前瞻性方法。 这不需要更改 XLA。 较小的小批量大小与 XLA 结合使用仍然可以优于原生 TensorFlow。

TensorFlow 流执行器并行执行许多操作。 这增加了内存需求，因为更多中间张量同时处于活动状态。 限制流执行器中的并行度通常会减少所需的内存量。 并行度由指定执行器可以使用的线程数来控制，方法是设置以下环境变量

TF_NUM_INTEROP_THREADS=<n>

9.3.2.1. 内存碎片

默认情况下，XLA 在一次分配中为中间张量分配所需的内存。 由于内存碎片，可能会发生 TensorFlow 分配器无法分配连续的内存缓冲区，即使仍有足够的内存可用。

通过设置

XLA_FLAGS=--xla_multiheap_size_constraint_per_heap=<n>

用户可以将单片内存分配分解为多个分配，每个分配的最大大小为 n 字节。

或者，TensorFlow-XLA 集成问题 描述了一个有时可以成功规避内存碎片问题的选项。

9.3.3. TensorFlow-XLA 性能问题

9.3.3.1. TensorFlow-XLA 集成问题

当 XLA 的性能比原生 TensorFlow 更差时，首先要尝试的事情之一是使用以下命令再次运行它

TF_XLA_FLAGS=--tf_xla_always_defer_compilation=true

如果性能仍然更差，则 TF-XLA 集成是问题所在，因为此选项将运行回退路径，而无需任何编译。 在这种情况下，使用以下命令再次运行

TF_XLA_FLAGS=--tf_xla_enable_lazy_compilation=false

这将完全更改 TensorFlow 与 XLA 接口的方式。 与集群中的图表不同，现在的图看起来像

请注意，回退路径不再存在。 此选项通常为给定形状实例多次执行的节点提供良好的性能结果。 它避免了合并节点和编译启发式，因为它必须在首次执行时进行编译。 此选项的直接副作用是非常不同的 graphDef 节点执行（执行单个 _XlaRun 节点而不是回退路径中的多个节点），因此 TensorFlow 内存分配器的使用方式非常不同。 在某些情况下，此选项成功避免了之前提到的内存碎片问题。 这仅仅是一个偶然的副作用，而不是修复。

TX-XLA 接口可能表现为性能问题的第二种方式是当出现许多较小的集群时。 当发生这种情况时，图中额外节点的开销、编译时间（尽管很短）和 XLA 执行器比 XLA 编译器实现的性能增益更差。 默认情况下，集群大小的下限为 4 个操作。 当使用以下命令运行时

TF_XLA_FLAGS=--tf_xla_min_cluster_size=<n>

其中 <n> 大于 4，任何操作少于 <n> 的集群都将被忽略，并且图保持不变。

9.3.3.2. 编译开销

编译时间开销可能会导致严重的性能下降，尤其是在输入数据形状不断变化的模型中。

要了解在编译上花费了多少时间，请使用以下命令运行模型

TF_CPP_VMODULE=xla_compilation_cache=1

这会在每次编译发生后转储信息。 它显示了上次编译花费了多长时间，以及到目前为止所有编译的累积时间。 当使用以下命令运行时

TF_CPP_VMODULE=nvptx_compiler=2

它还会转储有关 .ptx 文件和 .cub 文件大小的信息。 可以通过增加集群的下限（当存在许多小型编译时）来减少编译时间开销

TF_XLA_FLAGS=--tf_xla_min_cluster_size=<n>

或者通过减小集群大小的上限（当某些编译花费的时间过长时）

TF_XLA_FLAGS=--tf_xla_max_cluster_size=<n>

XLA 还可以在后台执行集群编译，而回退路径的执行可以取得进展。 这些后台编译会占用主执行管道的 CPU 资源，但在许多情况下，在等待编译完成时不阻塞执行可以降低端到端延迟。 可以通过设置选择加入异步编译

TF_XLA_FLAGS=--tf_xla_async_compilation=true

异步编译与禁用延迟编译不兼容。 当延迟编译被禁用时，它优先于异步编译。

XLA 可以将集群编译的部分内容持久缓存到磁盘。 此缓存可用于通过重用缓存的编译结果来加快后续运行的编译时间。 通过设置启用持久缓存

TF_XLA_FLAGS=--xla_gpu_persistent_cache_dir=<dir_name>

dirname 必须引用具有读取和写入权限的现有目录。 它将用于存储新条目和检索现有条目。 建议的使用方法是每个模型使用一个小缓存，因为所有缓存条目都在初始化时读取，因此模型未使用的条目将被读取，但永远不会使用。 这是一项实验性功能，不保证缓存与 TensorFlow 版本的兼容性。

9.3.3.3. 计算/通信重叠

如前所述，作为 XLA 如何与 TensorFlow 集成的直接副作用，_XlaRun 节点的所有输出都同时准备就绪。

这意味着任何 XLA 输出的消费者都必须等到生成所有输出。 这限制了 TensorFlow 执行器并行执行计算和通信。 事实是，在 XLA 集群内部，通常输出已经计算出来，并且可以传递给消费者。 XLA 有一个称为 AsyncIO 的优化，它允许 XLA 输出的消费者在输出可用时立即开始执行。

如果 XLA 在单个 GPU 上的性能优于原生 TensorFlow，但在多 GPU 系统上的性能更差，则可能是由于缺少计算和计算重叠。 使用以下命令运行 TensorFlow

TF_XLA_FLAGS=--tf_xla_async_io_level=1

启用 AsyncIO。

9.3.4. XLA 优化器和代码生成

在大多数情况下，当出现以下情况时，TF-XLA 接口问题会消失

• 运行足够长的时间。 随着集群的大多数形状实例在缓存中命中，编译开销会随着时间的推移而减少。

• 运行足够大的张量，即足够大的批量大小。

• 过滤掉微小的集群。

如果 XLA 编译的集群比回退路径运行得慢，则 XLA 编译器可能是性能下降的原因。

9.3.4.1. XLA 自动调优

默认情况下，XLA 对 cuBLAS 和 cuDNN 中的 GEMM 和 CONV 算法执行自动调优。 对于每个 GEMM，都会尝试 cuBLAS 中的所有可能的替代方案，并选择最快的方案； 对于 cuDNN 中的每个 CONV 操作也是如此。 自动调优由级别 [0..4] 控制。 级别 4 是默认值，它在不同的替代方案之间执行一些功能测试，以捕获 cuBLAS 和 cuDNN 问题。

尝试使用以下命令运行

XLA_FLAGS=--xla_gpu_autotune_level=2

这将执行完整的自动调优，但会跳过级别 4 中完成的一些初始化和比较。 对于短时间运行的脚本，甚至完全禁用自动调优也可以提高性能

XLA_FLAGS=--xla_gpu_autotune_level=0

9.3.4.2. 各种选项

cuDNN 批处理归一化

默认情况下，XLA 将批处理归一化层分解为许多较小的操作，然后由 XLA 优化这些操作。 仅对于训练，XLA 有一个选项可以将批处理归一化保留为单个操作，该操作以 cuDNN 批处理归一化 API 为目标来执行它。 使用以下命令运行

XLA_FLAGS=--xla_gpu_use_cudnn_batchnorm_level=2

以使用 cuDNN 执行前向和后向层的批处理归一化层。

tf.enable_resource_variables()

TensorFlow 资源变量 是 TensorFlow 变量的改进版本。 XLA 不会对变量进行集群，但会对资源变量进行集群。 因此，集群资源变量会集群更多的 graphDef 节点，包括 assign 节点。 我们至少看到一个模型，其中许多 assign 节点都依赖于 XLA 集群的输出。 启用资源变量使这些节点成为 XLA 集群的一部分，从而提高了延迟，因为这些 assign 节点不必等待 XLA 集群完成。 通过添加以下内容来实现启用资源变量

tf.enable_resource_variables()

在创建任何变量之前添加到模型源代码中。

--xla_backend_optimization_level

XLA_FLAGS=--xla_backend_optimization_level=0

--xla_gpu_disable_ptxas_optimizations

XLA_FLAGS=--xla_gpu_disable_ptxas_optimizations=true

9.4. XLA 选项参考

本节列出了本文档中提到的选项，并附有简短说明。

TF_XLA 选项

• --tf_xla_auto_jit

  控制集群。 0：关闭，1：开启，fusible：Xla-Lite

• --tf_xla_always_defer_compilation

  控制延迟编译。 true：回退路径，false：默认策略

• --tf_xla_max_cluster_size

  设置集群大小的上限

• --tf_xla_min_cluster_size

  设置集群大小的下限

• --tf_xla_enable_lazy_compilation

  控制延迟编译。 false：始终编译，true：默认策略

• --tf_xla_async_compilation

  控制异步编译。 false：默认策略，true：异步编译

  自 20.10 版本起可用

• --xla_gpu_persistent_cache_dir

  控制持久编译缓存。 当设置为现有目录时，该目录用作编译缓存。

  自 21.04 版本起可用

• --tf_xla_async_io_level

  启用异步 IO。 0：关闭，1：开启

XLA 选项

• --xla_gpu_autotune_level

  控制自动调优级别，0：关闭，1：无初始化，2：带初始化，3：带重新初始化，4：带功能检查

• --xla_multiheap_size_constraint_per_heap

  控制是否使用多个堆。 -1：单个堆，n：使用多个分配，每个分配的最大大小为 n 字节。

  自 20.11 版本起可用

• --xla_gpu_use_cudnn_batchnorm+level

  控制是否使用 cuDNN batchnorm。 0：否，1：仅 BatchNorm+Bias+Act，2：所有情况。

10.1. 支持

有关 TensorFlow 的更多信息，包括教程、文档和示例，请参阅：

• TensorFlow 教程
• TensorFlow API
• 自动混合精度
• Tensor Core
• 针对 Tensor Core 优化的深度学习示例

有关最新的 TensorFlow 发行说明，请参阅 https://docs.nvda.net.cn/deeplearning/frameworks/tensorflow-release-notes/index.html。

声明

本文档仅供参考，不得视为对产品的特定功能、条件或质量的保证。 NVIDIA Corporation (“NVIDIA”) 对本文档中包含的信息的准确性或完整性不做任何明示或暗示的陈述或保证，并且对本文档中包含的任何错误不承担任何责任。 NVIDIA 对因使用此类信息或因使用此类信息而可能导致的侵犯第三方专利或其他权利的后果或使用不承担任何责任。 本文档不构成开发、发布或交付任何材料（以下定义）、代码或功能的承诺。

NVIDIA 保留随时对此文档进行更正、修改、增强、改进和任何其他更改的权利，恕不另行通知。

客户在下订单之前应获取最新的相关信息，并应验证此类信息是否为最新且完整。

NVIDIA 产品的销售受 NVIDIA 订单确认时提供的标准销售条款和条件的约束，除非 NVIDIA 和客户的授权代表签署的个人销售协议（“销售条款”）另有约定。 NVIDIA 在此明确反对将任何客户通用条款和条件应用于购买本文档中引用的 NVIDIA 产品。 本文档未直接或间接形成任何合同义务。

NVIDIA 产品并非设计、授权或保证适用于医疗、军事、航空、航天或生命维持设备，也不适用于 NVIDIA 产品的故障或失灵可能合理预期会导致人身伤害、死亡或财产或环境损害的应用。 NVIDIA 对在上述设备或应用中包含和/或使用 NVIDIA 产品不承担任何责任，因此，此类包含和/或使用由客户自行承担风险。

NVIDIA 不保证基于本文档的产品适用于任何特定用途。 NVIDIA 不一定会对每个产品的所有参数进行测试。 客户全权负责评估和确定本文档中包含的任何信息的适用性，确保产品适合并满足客户计划的应用，并为应用执行必要的测试，以避免应用或产品的默认设置。 客户产品设计中的缺陷可能会影响 NVIDIA 产品的质量和可靠性，并可能导致超出本文档中包含的附加或不同条件和/或要求。 NVIDIA 对可能基于或归因于以下原因的任何默认设置、损坏、成本或问题不承担任何责任：(i) 以任何与本文档相悖的方式使用 NVIDIA 产品，或 (ii) 客户产品设计。

本文档未授予 NVIDIA 专利权、版权或本文档项下的其他 NVIDIA 知识产权的任何明示或暗示的许可。 NVIDIA 发布的有关第三方产品或服务的信息不构成 NVIDIA 授予使用此类产品或服务的许可，也不构成 NVIDIA 对其的保证或认可。 使用此类信息可能需要获得第三方在第三方的专利或其他知识产权项下的许可，或获得 NVIDIA 在 NVIDIA 的专利或其他知识产权项下的许可。

仅当事先获得 NVIDIA 的书面批准，在不更改且完全符合所有适用的出口法律和法规的情况下复制本文档中的信息，并附带所有相关的条件、限制和声明时，才允许复制本文档中的信息。

本文档和所有 NVIDIA 设计规范、参考板、文件、图纸、诊断程序、列表和其他文档（统称为“材料”）均按“原样”提供。 NVIDIA 不对材料作出任何明示、暗示、法定或其他方面的保证，并且明确声明不承担所有关于不侵权、适销性和特定用途适用性的暗示保证。 在法律未禁止的范围内，在任何情况下，NVIDIA 均不对任何损害（包括但不限于任何直接、间接、特殊、偶然、惩罚性或后果性损害）负责，无论其原因以及责任理论如何，即使 NVIDIA 已被告知此类损害的可能性，这些损害均因使用本文档而引起。 尽管客户可能因任何原因遭受任何损害，但 NVIDIA 对本文所述产品的总累积责任应根据产品的销售条款进行限制。

HDMI

HDMI、HDMI 徽标和 High-Definition Multimedia Interface 是 HDMI Licensing LLC 的商标或注册商标。

OpenCL

OpenCL 是 Apple Inc. 的商标，已获得 Khronos Group Inc. 的许可使用。

商标

NVIDIA、NVIDIA 徽标以及 cuBLAS、CUDA、DALI、DGX、DGX-1、DGX-2、DGX Station、DLProf、Jetson、Kepler、Maxwell、NCCL、Nsight Compute、Nsight Systems、NvCaffe、PerfWorks、Pascal、SDK Manager、Tegra、TensorRT、Triton Inference Server、Tesla、TF-TRT 和 Volta 是 NVIDIA Corporation 在美国和其他国家/地区的商标和/或注册商标。 其他公司和产品名称可能是与其相关的各自公司的商标。