理解为什么 EagerTensor.numpy() 很慢

Question

我正在尝试使用预先训练的自定义图像分割模型对某些图像数据进行预测，但我遇到了预测速度太慢的问题。使用Pyinstrument进行一些分析后，我发现对的调用numpy()占用了大部分运行时间。这让我感到惊讶，因为据我了解，该函数所需要做的就是将预测数据从 GPU 复制回主机，并可能执行一些其他操作将数据转换为 numpy 数组。

我制作了一个最小的工作示例来显示效果：

import numpy as np
import tensorflow as tf
from pyinstrument import Profiler

model = tf.keras.applications.resnet50.ResNet50(include_top=False, input_shape=(672, 2048, 3))
data = np.zeros((1, 672, 2048, 3))

profiler = Profiler()
profiler.start()

for _ in range(1000):
    prediction = model.predict_on_batch(data)
    output = prediction.numpy()

profiler.stop()
print(profiler.output_text())

该代码产生以下 Pyinstrument 输出：

  _     ._   __/__   _ _  _  _ _/_   Recorded: 15:50:54  Samples:  6759
 /_//_/// /_\ / //_// / //_'/ //     Duration: 56.536    CPU time: 56.406
/   _/                      v3.1.0

Program: test.py

56.536 <module>  test.py:1
|- 42.315 numpy  tensorflow_core\python\framework\ops.py:918
|     [4 frames hidden]  tensorflow_core
|        39.179 _numpy  tensorflow_core\python\framework\ops.py:905
|- 12.667 predict_on_batch  tensorflow_core\python\keras\engine\training.py:1220
|     [7705 frames hidden]  tensorflow_core, ast, _weakrefset, we...
`- 1.553 [self]

我在 RTX 2080 Ti 上使用 Tensorflow 2.1.0。当使用我自己的模型时，效果更加明显，在这种情况下，调用占用numpy()了超过 90% 的运行时间。

我也尝试过使用此评论output = np.array(memoryview(prediction))中提到的方法 ，但性能没有显着差异。

有人可以解释为什么会发生这种情况吗？这是预期的行为吗？调用 后输出是否未立即完全准备好predict_on_batch()，因此numpy()必须等待输出准备好才能从中创建数组？

更新：通过在示例代码的time.sleep(0.1)调用之间添加predict_on_batch()和 ，我得到以下输出：numpy()

  _     ._   __/__   _ _  _  _ _/_   Recorded: 12:03:52  Samples:  7821
 /_//_/// /_\ / //_// / //_'/ //     Duration: 135.388   CPU time: 25.641
/   _/                      v3.1.0

Program: test.py

135.388 <module>  test.py:1
|- 105.639 [self]
|- 16.191 numpy  tensorflow_core\python\framework\ops.py:918
|     [4 frames hidden]  tensorflow_core
`- 13.558 predict_on_batch  tensorflow_core\python\keras\engine\training.py:1220
      [7839 frames hidden]  tensorflow_core, ast, _weakrefset, we...

正如您所看到的， 的时间predict_on_batch()仍然大致相同，但是稍微等待一下， 花费的时间numpy()就大大减少了。对我来说，这似乎指向我的理论，即输出在到达时尚未完全准备好，numpy()因此它必须等待一段时间才能生成数组。但我仍然不太满意，有人可以解释一下这种情况吗？有什么办法可以提高性能吗？