Question

我是Tensorflow的新手。我有一个在CPU上运行的Tensorflow代码，但我试图将其更改为在GPU上运行。我在某处添加了tf.device('/device:GPU:0')语句，但是显然它不起作用，并且我也没有看到日志。

如何更改？

#...some stuff
g_in = tf.Graph()
with gfile.FastGFile(path + name + '.pb', 'rb') as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())
    tf.import_graph_def(graph_def, name="")
    g_in = tf.get_default_graph()
    print(g_in)

with g_in.as_default(), tf.device('/device:GPU:0'):
    sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
    input = sess.graph.get_tensor_by_name("cnn/input_1:0")
    output = sess.graph.get_tensor_by_name("mymodel:0")

Answer 1

如果您安装了tensorflow-gpu，默认情况下Tensorflow将在GPU上运行。
with tf.device('/device:GPU:0')用于选择要使用默认值的GPU。 pip freeze应该向您显示您具有哪种安装。如果仅显示tensorflow而不显示tensorflow-gpu，则需要卸载它，安装CUDA并重新安装tensorflow-gpu

在GPU上运行Tensorflow模型的问题

1 个答案: