Question

我编写了一个函数，使用带有tensorflow作为后端的keras，使用vgg16网络提取特征。问题是它默认在CPU而不是GPU上运行它。我确实有一台兼容GPU的机器，前几天另一个代码（用于培训）正在使用GPU。我添加了这段代码片段以强制它到GPU，但它仍然不起作用。

import tensorflow as tf
from keras import backend as K
GPU = True
CPU = False
num_cores = 4

if GPU:
    num_GPU = 1
    num_CPU = 1
if CPU:
    num_CPU = 1
    num_GPU = 0

config = tf.ConfigProto(intra_op_parallelism_threads=num_cores,\
    inter_op_parallelism_threads=num_cores, allow_soft_placement=True,\
    device_count = {'CPU' : num_CPU, 'GPU' : num_GPU})
session = tf.Session(config=config)
K.set_session(session)

from tensorflow.python.client import device_lib
print(device_lib.list_local_devices())}

这是运行特征提取后提供的输出

Backend Qt5Agg is interactive backend. Turning interactive mode on.
Using TensorFlow backend.
[name: "/cpu:0"
device_type: "CPU"
memory_limit: 268435456
locality {
}
incarnation: 5806448485889010842
]

以下是提取功能的代码（运行正常但在CPU上运行）

modelname  = 'vgg16'
Network    = MODELS[modelname]
model      = Network(weights="imagenet", include_top=False, input_shape=inputShape)
x          = preprocess(img_4D.copy())
features   = model.predict(np.float64(x))

Answer 1

您是否尝试卸载Tensorflow的CPU版本，以便计算机上唯一的版本是GPU版本？

使用Tensorflow后端强制Keras在GPU上运行

1 个答案: