python - 如何在 TF Keras 模型中测量图层执行的时间？

Question

我正在尝试找到一种方法来有效且舒适地测量 TF Keras 模型中单层的执行时间。我已经搜索了很多，但我还没有找到一个完全让我满意的解决方案。

我遇到的一件类似的事情是 Pytorch 中使用的一种方法，它基本上使用了时间库。对于 GPU 案例，它还使用了：torch.cuda.synchronize() 函数。

我可以在 TF 中做类似的事情吗？

我尝试过以下方式：

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)

但恐怕它不会是 GPU 的正确选择。至少在 CPU 上计算可以吗？

我还在 Tensorflow 中找到了有关 Eager Execution 模式的信息以及此类方法：

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

有人会推荐这种方法吗？

最后一件事，最常推荐的是 TF Profiler，但现在我觉得它对我的任务有点不舒服（可能是因为我不太了解它）。

最后，我想要实现的是拥有一个模型，在该模型中我可以迭代层，更改 Conv 操作中的过滤器数量，并根据输入和输出通道的数量计算单层的延迟。我将不胜感激任何想法！

标签： pythontensorflowkerastimegpu