如何使用 tf.keras 获取每个 train_step 的详细运行时间？

【问题标题】：How to get the detailed runtime per train_step using tf.keras?如何使用 tf.keras 获取每个 train_step 的详细运行时间？
【发布时间】：2020-06-03 14:37:15
【问题描述】：

训练时如何通过 tf.keras 获取每个 train_step 的详细运行时间？虽然我知道在 model.fit() 中添加回调能够返回纪元时间，但我需要知道：预测时间 loss_calculate_time 毕业时间毕业申请时间更新时间上面这样的操作就像我们使用这样的自定义训练循环所做的那样：

  def train_step(inputs):
    images, labels = inputs

    with tf.GradientTape() as tape:
      predictions = model(images, training=True)
      loss = compute_loss(labels, predictions)

    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    train_accuracy.update_state(labels, predictions)
    return loss

【问题讨论】：

您可以通过回调获得损失和指标。您无法获得任何执行时间。但是您可以为“批处理执行时间”编写回调（不低于此）。

标签： python tensorflow keras tensorflow2.0

【解决方案1】：

使用Profiler 提供的Tensorflow 功能可以满足您的要求。

Training Step 期间每个操作的统计信息可以使用Tensorflow Profiler 捕获，并且可以在Tensorboard 的Profile 选项卡中可视化。

由于您有兴趣了解在训练期间发生的不同操作的执行时间，因此可以在Profiler 的 tensorflow_stats 下拉菜单中对其进行可视化，如下面的屏幕截图所示。

可以使用Tensorboard 中的Profiler 可视化的示例操作的屏幕截图（64 个屏幕截图中只有 17 个可见）如下所示：

如需了解更多有关此功能的信息，Tensorflow Profiler，请参考this Guide、this Tutorial 和此Dev Summit Youtube Video。

希望这会有所帮助。学习愉快！！

【讨论】：