如何在 Tensorflow 的深度学习期间测量每批次的训练时间？

Question

I want to measure training time per batches during Deep Learning in Tensorflow. There are several ways to measure training time per epochs, but I cannot find how to measure training time per batches.

I tried Tensorboard, but I don't know how to add some kind of 'execution time' scalars in tensorboard callbacks.

And I also tried to override the function 'train_step' but it didn't work the way I want.

Answer 1

You can do this by creating a keras custom callback. The code below will print out the time it takes to process each batch, the training accuracy for that batch and the loss for that batch

class batch_timer(keras.callbacks.Callback):
    def __init__(self  ):
        super(batch_timer, self).__init__()
        
    def on_train_batch_begin(self, batch, logs=None):
        self.start_time=time.time()    
    def on_train_batch_end(self, batch, logs=None):
        stop_time=time.time()
        duration =stop_time-self.start_time
        acc=logs.get('accuracy')* 100  # get training accuracy 
        loss=logs.get('loss')
        msg='processing batch {0:6s}  duration= {1:12.4f}  accuracy= {2:8.3f}  loss: {3:8.5f}'.format(str(batch), duration,  acc, loss)
        print(msg ) 

Now you are going to get tons of data. When you run model.fit include

history=model.fit(etc.... ,verbose=0, callbacks=[batch_timer()]

I set verbose=0 because if you don't the printout gets a bit messy because model.fit printing gets messed up by the callback printing