首页 > 解决方案 > 在使用 keras Resnet50 进行预测时使用 gpu

问题描述

我正在使用 Resnet50 模型对图像进行矢量化以查找图像相似性。为了提高程序的速度,我尝试了多处理,但由于 keras 的后端逻辑而失败了。最后,我在单机上运行我的矢量化代码,它的性能还不错,但我需要更好。为了减少这个向量化操作所消耗的时间,我可以使用我的 gpu 包含的机器。但是我在调​​用预测方法时找不到使用 gpu 的方法。任何帮助对我来说都会很棒。

示例代码:

basemodel = tensorflow.keras.applications.ResNet50(weights='imagenet', include_top=False, pooling="avg", input_shape=self.input_shape) 
model = tensorflow.keras.models.Model(inputs=basemodel.input, outputs=basemodel.output)
img_data = img_to_array(image)
img_data = np.expand_dims(img_data, axis=0)
img_data = preprocess_input(img_data)
feature_vector = basemodel.predict(img_data)

我需要加快速度basemodel.predict(img_data)。我可以为此目的使用 gpu 吗?

标签: image-processinggputpu

解决方案


由于您从 tensorflow.keras 导入了模型,因此您无需更改代码即可使用 GPU。您应该只具备先决条件:https ://www.tensorflow.org/install/gpu#software_requirements 。您可以检查您的程序是否看到具有以下内容的 gpu 设备:

import tensorflow as tf;

print("Num GPUs Available: ", len(tf.config.experimental.list_physical_devices('GPU')))

nvidia-smi最后,您可以在终端上运行命令检查 gpu 利用率


推荐阅读