首页 > 解决方案 > 在 GPU 上运行我的 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")

标签: python-3.xtensorflowmachine-learning

解决方案


tensorflow-gpu如果您已安装,Tensorflow 将默认在 GPU 上运行。
with tf.device('/device:GPU:0')如果您不想在默认情况下运行它,则用于选择要使用的 GPU。 pip freeze应该告诉你你有哪个安装。如果它只说tensorflow而不是tensorflow-gpu那么你需要卸载它,安装 CUDA 并重新安装tensorflow-gpu


推荐阅读