#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Explanation: Copyright 2019 The TensorFlow Authors.
End of explanation
"""


try:
  # %tensorflow_version 只在 Colab 中存在。
  %tensorflow_version 2.x
except Exception:
  pass

# 加载 TensorBoard notebook 扩展。
%load_ext tensorboard


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from datetime import datetime
from packaging import version

import functools
import tensorflow as tf
import tensorflow_datasets as tfds
from tensorflow.python.keras import backend
from tensorflow.python.keras import layers

import numpy as np

print("TensorFlow version: ", tf.__version__)


"""
Explanation: TensorBoard 性能分析: 在 Keras 中对基本训练指标进行性能分析
<table class="tfo-notebook-buttons" align="left">
  <td>
    <a target="_blank" href="https://tensorflow.google.cn/tensorboard/tensorboard_profiling_keras"><img src="https://tensorflow.google.cn/images/tf_logo_32px.png" />在 TensorFlow.org 上查看</a>
  </td>
  <td>
    <a target="_blank" href="https://colab.research.google.com/github/tensorflow/docs-l10n/blob/master/site/zh-cn/tensorboard/tensorboard_profiling_keras.ipynb"><img src="https://tensorflow.google.cn/images/colab_logo_32px.png" />在 Google Colab 上运行</a>
  </td>
  <td>
    <a target="_blank" href="https://github.com/tensorflow/docs-l10n/blob/master/site/zh-cn/tensorboard/tensorboard_profiling_keras.ipynb"><img src="https://tensorflow.google.cn/images/GitHub-Mark-32px.png" />在 GitHub 上查看源代码</a>
  </td>
  <td>
    <a href="https://storage.googleapis.com/tensorflow_docs/docs-l10n/site/zh-cn/tensorboard/tensorboard_profiling_keras.ipynb"><img src="https://tensorflow.google.cn/images/download_logo_32px.png" />下载此 notebook</a>
  </td>
</table>

总览
在机器学习中性能十分重要。TensorFlow 有一个内置的性能分析器可以使您不用费力记录每个操作的运行时间。然后您就可以在 TensorBoard 的 Profile Plugin 中对配置结果进行可视化。本教程侧重于 GPU ，但性能分析插件也可以按照云 TPU 工具来在 TPU 上使用。
本教程提供了非常基础的示例以帮助您学习如何在开发 Keras 模型时启用性能分析器。您将学习如何使用 Keras TensorBoard 回调函数来可视化性能分析结果。“其他性能分析方式”中提到的 Profiler API 和 Profiler Server 允许您分析非 Keras TensorFlow 的任务。
事先准备


在你的本地机器上安装最新的TensorBoard。


在 Notebook 设置的加速器的下拉菜单中选择 “GPU”（假设您在Colab上运行此notebook）





设置
End of explanation
"""


device_name = tf.test.gpu_device_name()
if not tf.test.is_gpu_available():
  raise SystemError('GPU device not found')
print('Found GPU at: {}'.format(device_name))

"""
Explanation: 确认 TensorFlow 可以看到 GPU。
End of explanation
"""


BATCH_NORM_DECAY = 0.997
BATCH_NORM_EPSILON = 1e-5
L2_WEIGHT_DECAY = 2e-4


def identity_building_block(input_tensor,
                            kernel_size,
                            filters,
                            stage,
                            block,
                            training=None):
 
  """标识块是一种在捷径上没有卷积层的块。

  参数：
    input_tensor：输入张量
    kernel_size：默认为3，内核大小为
        主路径上的中间卷积层
    过滤器：整数列表，主路径上3个卷积层的过滤器
    stage：整数，当前阶段标签，用于生成层名称
    block：当前块标签，用于生成层名称
    training：仅在使用 Estimator 训练 keras 模型时使用。 在其他情况下，它是自动处理的。

  返回值：
    输出块的张量。
  """
  filters1, filters2 = filters
  if tf.keras.backend.image_data_format() == 'channels_last':
    bn_axis = 3
  else:
    bn_axis = 1
  conv_name_base = 'res' + str(stage) + block + '_branch'
  bn_name_base = 'bn' + str(stage) + block + '_branch'

  x = tf.keras.layers.Conv2D(filters1, kernel_size,
                             padding='same',
                             kernel_initializer='he_normal',
                             kernel_regularizer=
                             tf.keras.regularizers.l2(L2_WEIGHT_DECAY),
                             bias_regularizer=
                             tf.keras.regularizers.l2(L2_WEIGHT_DECAY),
                             name=conv_name_base + '2a')(input_tensor)
  x = tf.keras.layers.BatchNormalization(axis=bn_axis,
                                         name=bn_name_base + '2a',
                                         momentum=BATCH_NORM_DECAY,
                                         epsilon=BATCH_NORM_EPSILON)(
                                             x, training=training)
  x = tf.keras.layers.Activation('relu')(x)

  x = tf.keras.layers.Conv2D(filters2, kernel_size,
                             padding='same',
                             kernel_initializer='he_normal',
                             kernel_regularizer=
                             tf.keras.regularizers.l2(L2_WEIGHT_DECAY),
                             bias_regularizer=
                             tf.keras.regularizers.l2(L2_WEIGHT_DECAY),
                             name=conv_name_base + '2b')(x)
  x = tf.keras.layers.BatchNormalization(axis=bn_axis,
                                         name=bn_name_base + '2b',
                                         momentum=BATCH_NORM_DECAY,
                                         epsilon=BATCH_NORM_EPSILON)(
                                             x, training=training)

  x = tf.keras.layers.add([x, input_tensor])
  x = tf.keras.layers.Activation('relu')(x)
  return x


def conv_building_block(input_tensor,
                        kernel_size,
                        filters,
                        stage,
                        block,
                        strides=(2, 2),
                        training=None):
  """在捷径中具有卷积层的块。

  参数：
    input_tensor：输入张量
    kernel_size：默认为3，内核大小为
        主路径上的中间卷积层
    filters：整数列表，主路径上3个卷积层的过滤器
    stage：整数，当前阶段标签，用于生成层名称
    block：当前块标签，用于生成层名称
    training：仅在使用 Estimator 训练 keras 模型时使用。在其他情况下，它是自动处理的。

  返回值：
    输出块的张量。

  请注意，从第3阶段开始，
  主路径上的第一个卷积层的步长=（2，2）
  而且捷径的步长=（2，2）
  """
  filters1, filters2 = filters
  if tf.keras.backend.image_data_format() == 'channels_last':
    bn_axis = 3
  else:
    bn_axis = 1
  conv_name_base = 'res' + str(stage) + block + '_branch'
  bn_name_base = 'bn' + str(stage) + block + '_branch'

  x = tf.keras.layers.Conv2D(filters1, kernel_size, strides=strides,
                             padding='same',
                             kernel_initializer='he_normal',
                             kernel_regularizer=
                             tf.keras.regularizers.l2(L2_WEIGHT_DECAY),
                             bias_regularizer=
                             tf.keras.regularizers.l2(L2_WEIGHT_DECAY),
                             name=conv_name_base + '2a')(input_tensor)
  x = tf.keras.layers.BatchNormalization(axis=bn_axis,
                                         name=bn_name_base + '2a',
                                         momentum=BATCH_NORM_DECAY,
                                         epsilon=BATCH_NORM_EPSILON)(
                                             x, training=training)
  x = tf.keras.layers.Activation('relu')(x)

  x = tf.keras.layers.Conv2D(filters2, kernel_size, padding='same',
                             kernel_initializer='he_normal',
                             kernel_regularizer=
                             tf.keras.regularizers.l2(L2_WEIGHT_DECAY),
                             bias_regularizer=
                             tf.keras.regularizers.l2(L2_WEIGHT_DECAY),
                             name=conv_name_base + '2b')(x)
  x = tf.keras.layers.BatchNormalization(axis=bn_axis,
                                         name=bn_name_base + '2b',
                                         momentum=BATCH_NORM_DECAY,
                                         epsilon=BATCH_NORM_EPSILON)(
                                             x, training=training)

  shortcut = tf.keras.layers.Conv2D(filters2, (1, 1), strides=strides,
                                    kernel_initializer='he_normal',
                                    kernel_regularizer=
                                    tf.keras.regularizers.l2(L2_WEIGHT_DECAY),
                                    bias_regularizer=
                                    tf.keras.regularizers.l2(L2_WEIGHT_DECAY),
                                    name=conv_name_base + '1')(input_tensor)
  shortcut = tf.keras.layers.BatchNormalization(
      axis=bn_axis, name=bn_name_base + '1',
      momentum=BATCH_NORM_DECAY, epsilon=BATCH_NORM_EPSILON)(
          shortcut, training=training)

  x = tf.keras.layers.add([x, shortcut])
  x = tf.keras.layers.Activation('relu')(x)
  return x


def resnet_block(input_tensor,
                 size,
                 kernel_size,
                 filters,
                 stage,
                 conv_strides=(2, 2),
                 training=None):
  """一个应用层后跟多个标识块的块。

  参数：
    input_tensor：输入张量
    size：整数，构成转化卷积/身份块的数量。
    一个卷积层使用后，再跟（size-1）个身份块。
    kernel_size：默认为3，内核大小为
        主路径上的中间卷积层
    filters：整数列表，主路径上3个卷积层的过滤器
    stage：整数，当前阶段标签，用于生成层名称
    conv_strides：块中第一个卷积层的步长。
    training：仅在使用 Estimator 训练 keras 模型时使用。其他情况它会自动处理。  

  返回值：
    应用层和身份块后的输出张量。
  """

  x = conv_building_block(input_tensor, kernel_size, filters, stage=stage,
                          strides=conv_strides, block='block_0',
                          training=training)
  for i in range(size - 1):
    x = identity_building_block(x, kernel_size, filters, stage=stage,
                                block='block_%d' % (i + 1), training=training)
  return x

def resnet(num_blocks, classes=10, training=None):
  """实例化ResNet体系结构。

  参数：
    num_blocks：整数，每个块中的卷积/身份块的数量。
      ResNet 包含3个块，每个块包含一个卷积块
      后面跟着(layers_per_block - 1) 个身份块数。 每
      卷积/理想度块具有2个卷积层。 用输入
      卷积层和池化层至最后，这带来了
      网络的总大小为（6 * num_blocks + 2）
    classes：将图像分类为的可选类数
    training：仅在使用 Estimator 训练 keras 模型时使用。其他情况下它会自动处理。

  返回值：
    Keras模型实例。
  """

  input_shape = (32, 32, 3)
  img_input = layers.Input(shape=input_shape)

  if backend.image_data_format() == 'channels_first':
    x = layers.Lambda(lambda x: backend.permute_dimensions(x, (0, 3, 1, 2)),
                      name='transpose')(img_input)
    bn_axis = 1
  else:  # channel_last
    x = img_input
    bn_axis = 3

  x = tf.keras.layers.ZeroPadding2D(padding=(1, 1), name='conv1_pad')(x)
  x = tf.keras.layers.Conv2D(16, (3, 3),
                             strides=(1, 1),
                             padding='valid',
                             kernel_initializer='he_normal',
                             kernel_regularizer=
                             tf.keras.regularizers.l2(L2_WEIGHT_DECAY),
                             bias_regularizer=
                             tf.keras.regularizers.l2(L2_WEIGHT_DECAY),
                             name='conv1')(x)
  x = tf.keras.layers.BatchNormalization(axis=bn_axis, name='bn_conv1',
                                         momentum=BATCH_NORM_DECAY,
                                         epsilon=BATCH_NORM_EPSILON)(
                                             x, training=training)
  x = tf.keras.layers.Activation('relu')(x)

  x = resnet_block(x, size=num_blocks, kernel_size=3, filters=[16, 16],
                   stage=2, conv_strides=(1, 1), training=training)

  x = resnet_block(x, size=num_blocks, kernel_size=3, filters=[32, 32],
                   stage=3, conv_strides=(2, 2), training=training)

  x = resnet_block(x, size=num_blocks, kernel_size=3, filters=[64, 64],
                   stage=4, conv_strides=(2, 2), training=training)

  x = tf.keras.layers.GlobalAveragePooling2D(name='avg_pool')(x)
  x = tf.keras.layers.Dense(classes, activation='softmax',
                            kernel_initializer='he_normal',
                            kernel_regularizer=
                            tf.keras.regularizers.l2(L2_WEIGHT_DECAY),
                            bias_regularizer=
                            tf.keras.regularizers.l2(L2_WEIGHT_DECAY),
                            name='fc10')(x)

  inputs = img_input
  # 创建模型
  model = tf.keras.models.Model(inputs, x, name='resnet56')

  return model


resnet20 = functools.partial(resnet, num_blocks=3)
resnet32 = functools.partial(resnet, num_blocks=5)
resnet56 = functools.partial(resnet, num_blocks=9)
resnet110 = functools.partial(resnet, num_blocks=18)

"""
Explanation: 使用 TensorBoard callback 运行一个简单的模型
你将使用 Keras 来构建一个使用 ResNet56 (参考: 用于图像识别的深度残差学习)来分类CIFAR-10图像集的简单模型。
从 TensorFlow 模型园复制 ResNet 模型代码。
End of explanation
"""


cifar_builder = tfds.builder('cifar10')
cifar_builder.download_and_prepare()


"""
Explanation: 从 TensorFlow 数据集下载 CIFAR-10 数据集。
End of explanation
"""


HEIGHT = 32
WIDTH = 32
NUM_CHANNELS = 3
NUM_CLASSES = 10
BATCH_SIZE = 128

def preprocess_data(record):
  image = record['image']
  label = record['label']
  
  # 调整图像大小以在每侧增加四个额外的像素。
  image = tf.image.resize_with_crop_or_pad(
      image, HEIGHT + 8, WIDTH + 8)

  # 随机裁剪图像的 [HEIGHT，WIDTH] 部分。
  image = tf.image.random_crop(image, [HEIGHT, WIDTH, NUM_CHANNELS])

  # 随机水平翻转图像。
  image = tf.image.random_flip_left_right(image)

  # 减去均值并除以像素方差。
  image = tf.image.per_image_standardization(image)
  
  label = tf.compat.v1.sparse_to_dense(label, (NUM_CLASSES,), 1)
  return image, label

train_data = cifar_builder.as_dataset(split=tfds.Split.TRAIN)
train_data = train_data.repeat()
train_data = train_data.map(
    lambda value: preprocess_data(value))
train_data = train_data.shuffle(1024)

train_data = train_data.batch(BATCH_SIZE)

model = resnet56(classes=NUM_CLASSES)

model.compile(optimizer='SGD',
              loss='categorical_crossentropy',
              metrics=['categorical_accuracy'])


"""
Explanation: 建立数据输入线性通信模型并编译 ResNet56 模型。
End of explanation
"""


log_dir="logs/profile/" + datetime.now().strftime("%Y%m%d-%H%M%S")

tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=log_dir, histogram_freq=1, profile_batch = 3)

"""
Explanation: 当你创建 TensorBoard 回调时，您可以指定您想要进行性能分析的批次。默认情况下，TensorFlow 将对第二个批次进行性能分析，因为第一个批次的时候会运行很多一次性的图优化。您可以通过设置 profile_batch 对其进行修改。您还可以通过将其设置为 0 来关闭性能分析。
这时候，您将会对第三批次进行性能分析。
End of explanation
"""


model.fit(train_data,
          steps_per_epoch=20,
          epochs=5, 
          callbacks=[tensorboard_callback])

"""
Explanation: 开始使用 Model.fit() 进行训练。
End of explanation
"""


!tar -zcvf logs.tar.gz logs/profile/

"""
Explanation: 使用 TensorBoard 可视化性能分析结果
不幸的是，由于＃1913, 您无法在 Colab 中使用 TensorBoard 来可视化性能分析结果。您需要下载日志目录并在本地计算机上启动 TensorBoard。
压缩下载日志:
End of explanation
"""


train_data = cifar_builder.as_dataset(split=tfds.Split.TRAIN)
train_data = train_data.repeat()
train_data = train_data.map(
    lambda value: preprocess_data(value))
train_data = train_data.shuffle(1024)
train_data = train_data.batch(BATCH_SIZE)

# 它将在（s-1）步骤中预取数据
train_data = train_data.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)

"""
Explanation: 在“文件”选项卡中右键单击以下载 logdir.tar.gz。

请保证在你本地的机器安装最新的 TensorBoard。在你的本地机器上执行下面的命令：
```

cd download/directory
tar -zxvf logs.tar.gz
tensorboard --logdir=logs/ --port=6006

```
在您的Chrome浏览器中打开一个新标签，然后导航至localhost：6006，单击 “Profile” 标签。您可能会看到以下性能分析结果：

跟踪查看器
当您单击性能分析选项卡后，您将看到跟踪查看器。该页面显示了聚合期间 CPU 和加速器上发生的不同事件的时间轴。
跟踪查看器在垂直轴上显示多个 事件组。 每个事件组都有多个水平 跟踪，其中填充了跟踪事件。跟踪 事件是在线程或 GPU 流上执行的基本时间线，。单个事件是时间轴轨道上的彩色矩形块。时间从左到右移动。
您可以使用 w（放大），s（缩小），a（向左滚动），d（向右滚动）浏览结果。
单个矩形代表 跟踪事件 ：从这个时间的开始到结束时间。 要研究单个矩形，可以在浮动工具栏中选择鼠标光标图标后单击它。 这将显示有关矩形的信息，例如其开始时间和持续时间。
除了点击之外，您还可以拖动鼠标以选择覆盖一组跟踪事件的矩形。这将为您提供与该矩形相交并汇总的事件列表。 m 键可用于测量所选事件的持续时间。

跟踪事件是从三个来源收集的：

CPU: CPU事件位于名为/host:CPU的事件组下。每个轨道代表 CPU 上的一个线程。例如，输入线性通信模型事件，GPU 操作调度事件， CPU 操作执行事件等。
GPU:  GPU 事件位于以 /device:GPU:为前缀的事件组下。 除了 stream:all，每个事件组都代表在 GPU 上一个流。 stream::all将所有事件汇集到一个 GPU 上。例如。 内存复制事件，内核执行事件等。
TensorFlow 运行时间: 运行时事件在以 /job:为前缀的事件组下。运行事件表示 python 程序调用的 TensorFlow ops。 例如， tf.function 执行事件等。

调试性能
现在，您将使用 Trace Viewer 来改善您的模型的性能。
让我们回到刚刚捕获的分析结果。

GPU 事件表明，GPU 在该步骤的上半部分什么都没有做。

CPU 事件表明，在此步骤的开始的时候，CPU 被数据输入管道占用。

在 TensorFlow 运行时中，有一个叫 Iterator::GetNextSync的大阻塞，这是从数据输入管道中获取下一批的阻塞调用。而且它阻碍了训练步骤。 因此，如果您可以在 s-1 的时候为 s 步骤准备输入数据，则可以更快地训练该模型。
您也可以通过使用 tf.data.prefetch.
End of explanation
"""


log_dir="logs/profile/" + datetime.now().strftime("%Y%m%d-%H%M%S")

tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=log_dir, histogram_freq=1, profile_batch = 3)

model.fit(train_data,
          steps_per_epoch=20,
          epochs=5, 
          callbacks=[tensorboard_callback])

"""
Explanation: 重新运行模型。
End of explanation
"""


!tar -zcvf logs.tar.gz logs/profile/

"""
Explanation: Woohoo! 你刚刚把训练性能从 ~235ms/step 提高到 ~200ms/step。
End of explanation
"""


# 内容管理接口
with tf.python.eager.profiler.Profiler('logdir_path'):
  # 进行你的训练
  pass


# 功能接口
tf.python.eager.profiler.start()
# 进行你的训练
profiler_result = tf.python.eager.profiler.stop()
tf.python.eager.profiler.save('logdir_path', profiler_result)

"""
Explanation: 再一次下载 logs 目录来查看 TensorBoard的新的分析结果。

Iterator::GetNextSync大阻塞不再存在。
做得好！
显然，这依旧不是最佳性能。请自己尝试，看看是否可以有更多的改进。
有关性能调整的一些有用参考：

数据输入线性通信模型
训练表现: 更快收敛的用户指南 (TensorFlow Dev Summit 2018)

其他分析方式
除了 TensorBoard 回调外，TensorFlow 还提供了其他两种方式来手动触发分析器：Profiler APIs 和 Profiler Service。
注意：请不要同时运行多个分析器。如果您想将 Profiler API 或 Profiler Service 与 TensorBoard 回调一起使用，请确保将profile_batch 参数设置为0。
Profiler APIs
End of explanation
"""


# 此 API 将在您的 TensorFlow 作业上启动 gRPC 服务器，该 API 可以按需接收分析请求。
tf.python.eager.profiler.start_profiler_server(6009)

# 在这里写你的 TensorFlow 项目

"""
Explanation: Profiler Service
End of explanation
"""


