小编Ost*_*ger的帖子

我正在尝试使用预先训练的自定义图像分割模型对某些图像数据进行预测，但我遇到了预测速度太慢的问题。使用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 …