我在训练时如何通过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
我找不到张量流各个操作所花费的时间的实现。但是,您可以使用on_batch_begin
的on_batch_end
,on_predict_batch_begin
,on_predict_batch_end
,tf.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.