jupyter notebook tensorflow打印device信息

最新推荐文章于 2024-11-26 22:09:46 发布

原创最新推荐文章于 2024-11-26 22:09:46 发布 · 998 阅读

2 ·

CC 4.0 BY-SA版权

tensorflow 专栏收录该内容

1 篇文章

订阅专栏

本文介绍如何在TensorFlow中设置设备分配并打印设备信息。通过使用log_device_placement参数和output_partition_graphs选项，可以在运行会话时获取到详细的设备分配情况。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

juypter notebook中直接使用log_device_placement=True打印不出来device信息

# Creates a graph.
with tf.device('/device:CPU: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,allow_soft_placement=True))
# Runs the op.

print(sess.run(c))

需要使用output_partition_graphs来输出device信息

# Creates a graph.
with tf.device('/device:GPU: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,allow_soft_placement=True))

# Runs the op.
options = tf.RunOptions(output_partition_graphs=True)
metadata = tf.RunMetadata()
c_val = sess.run(c, options=options, run_metadata=metadata)

print metadata.partition_graphs