首页 > 解决方案 > 我想检查 gpu 是否在 tensorflow 版本 1 中使用

问题描述

gpu 仍然低于 5%,是否有任何代码可以检查模型是否使用 gpu?(张量流 1.15)

标签: pythontensorflow

解决方案


使用下面的代码片段告诉 Tensorflow 使用 GPU

import tensorflow as tf 
if tf.test.gpu_device_name():
  print('Default GPU Device:{}'.format(tf.test.gpu_device_name()))

else:
  print("Please install GPU version of TF")

输出

Default GPU Device:/device:GPU:0

如果您希望特定操作在您选择的设备上运行,而不是自动为您选择的设备,您可以使用 withtf.device创建一个设备上下文,该上下文中的所有操作都将在同一指定设备上运行。

tf.debugging.set_log_device_placement(True)

# Place tensors on the CPU
with tf.device('/GPU:0'):
  a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
  b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])

# Run on the GPU
c = tf.matmul(a, b)
print(c)

推荐阅读