如何检查tensorflow环境是否能正常调用GPU

最新推荐文章于 2024-08-31 11:22:09 发布

转载最新推荐文章于 2024-08-31 11:22:09 发布 · 3.9k 阅读

1 ·

CC 4.0 BY-SA版权

原文链接：http://www.cnblogs.com/zmshare/p/10723820.html

文章标签：

#人工智能 #python #操作系统

本文通过在TensorFlow环境中创建并运行矩阵乘法操作，演示了如何验证GPU是否被正确调用。代码首先设置了环境变量以指定使用第0号GPU，并创建了两个常量张量，然后进行矩阵乘法运算。最后，通过设置会话配置参数log_device_placement为True，确保GPU被用于计算。

检查keras/tensorflow是否正常调用GPU代码

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    os.environ['CUDA_VISIBLE_DEVICES'] = "0"
    a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')
    b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b')
    c = tf.matmul(a, b)
    # Creates a session with log_device_placement set to True.
    sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
    # Runs the op.
    print(sess.run(c))
    # print(device_lib.list_local_devices())