I'm trying find a way to efficiently and comfortably measure time of execution of a single layer in TF Keras model. I've already searched quite a lot, but I haven't found a solution that fully satisfies me.
One thing I've came across, which was similar, was an approach used in Pytorch which basically used time library. And for GPU case, it also used: torch.cuda.synchronize() function.
Can I do something similar in TF?
I've tried in a following way:
def measure_layer_latency(layer, batch_size=4, runtimes=5):
input_shape = (batch_size,) + tuple(layer.input_shape[1:]) #input data shape including batch size
total_time = .0
for i in range(runtimes):
x = tf.random.normal(input_shape)
start = time.time()
layer(x)
finish = time.time()
total_time += (finish-start)
return total_time/float(runtimes)
but I'm afraid it won't be a right choice for GPU. Is it at least fine for calculations on CPU?
I've also found information about Eager Execution mode in Tensorflow and such method:
import time
def measure(x, steps):
# TensorFlow initializes a GPU the first time it's used, exclude from timing.
tf.matmul(x, x)
start = time.time()
for i in range(steps):
x = tf.matmul(x, x)
# tf.matmul can return before completing the matrix multiplication
# (e.g., can return after enqueing the operation on a CUDA stream).
# The x.numpy() call below will ensure that all enqueued operations
# have completed (and will also copy the result to host memory,
# so we're including a little more than just the matmul operation
# time).
_ = x.numpy()
end = time.time()
return end - start
Would anybody recommend such approach?
The last thing, most often recommended is TF Profiler, but for now I find it a bit uncomfortable for my task (maybe because I don't know it well).
Finally, the thing which I would like to achieve is having a model in which I can iterate over layers, change number of filters in Conv operations and calculate latency of a single layer depending on the number of input and output channels. I will be grateful for any ideas!