tensorflow如何用gpu训练?

1、确保tensorflow是gpu版本的

import tensorflow as tf
print("Num GPUs Available: ", len(tf.config.experimental.list_physical_devices('GPU')))

运行结果:Num GPUs Available: 2

 

2、TensorFlow支持在各种类型的设备上运行计算,包括CPU和GPU。 它们用字符串标识符表示,例如:

(1)"/device:CPU:0": The CPU of your machine.

(2)"/GPU:0": Short-hand notation for the first GPU of your machine that is visible to TensorFlow.

(3)"/job:localhost/replica:0/task:0/device:GPU:1": 你的机器的第二个可用的GPU的全称。

如果TensorFlow操作同时具有CPU和GPU实施,则默认情况下,将操作分配给设备时,GPU设备将获得优先级。 例如,tf.matmul同时具有CPU和GPU内核。 在具有CPU:0和GPU:0的设备的系统上,除非您明确要求在另一台设备上运行,否则将选择GPU:0设备运行tf.matmul。

3、记录操作分配给哪些设备

要找出操作和张量分配给哪些设备,请将tf.debugging.set_log_device_placement(True)作为程序的第一条语句。 启用设备放置日志记录将导致打印任何张量分配或操作。

4、手动指定设备

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

tf.debugging.set_log_device_placement(True)
 
# Place tensors on the CPU
with tf.device('/CPU: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]])
 
c = tf.matmul(a, b)
print(c)

5、限制GPU内存增长

默认情况下,TensorFlow会映射该进程可见的几乎所有GPU的所有GPU内存(取决于CUDA_VISIBLE_DEVICES)。 这样做是为了通过减少内存碎片来更有效地使用设备上相对宝贵的GPU内存资源。 要将TensorFlow限制为一组特定的GPU,我们使用tf.config.experimental.set_visible_devices方法。

原文链接:https://blog.csdn.net/qq_31554953/article/details/107302404

posted @ 2021-10-13 16:00  Tomorrow1126  阅读(487)  评论(0编辑  收藏  举报