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

问题描述 投票:1回答:1

我在训练时如何通过tf.keras获得每个train_step的详细运行时间?虽然我知道将回调添加到model.fit()能够返回时代,但我需要知道:预测时间loss_calculate_timegrad_timegrad_apply_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个回答
0
投票

我找不到张量流各个操作所花费的时间的实现。但是,您可以使用on_batch_beginon_batch_endon_predict_batch_beginon_predict_batch_endtf.keras.callbacks.Callback来计算一些需求,例如训练时间,批处理时间或预测时间。回调中还有更多选项,您可以找到它们here

捕获时间的示例代码-

import time
start = time.perf_counter()

def myfunc():
  a = 0
  for i in range(10000):
    a = a + 1
  return a

b = myfunc()

elapsed = time.perf_counter() - start

print('Elapsed time for myfunc() is %.3f seconds.' % elapsed)
# The .3f is to round to 3 decimal places.

输出-

Elapsed time for myfunc() is 0.001 seconds.
© www.soinside.com 2019 - 2024. All rights reserved.