TensorFlow使用GPU(tf.device)

本文介绍如何在TensorFlow中管理计算设备,包括CPU和GPU。通过tf.device函数可以显式指定运算在哪个设备上执行,同时利用log_device_placement参数可以查看每个运算的实际执行设备。
部署运行你感兴趣的模型镜像

        TensorFlow程序可以通过tf.device函数来指定运行每一个操作的设备,这个设备可以是本地的CPU或者GPU,也可以是某一台远程的服务器。TensorFlow会给每一个可用的设备一个名称,tf.device函数可以通过设备的名称来指定执行运算的设备。比如CPU在TensorFlow中的名称为/cpu:0。

在默认情况下,即使机器有多个CPU,TensorFlow也不会区分它们,所有的CPU都使用/cpu:0作为名称。而一台机器上不同GPU的名称是不同的,第n个GPU在TensorFlow中的名称为/gpu:n。比如第一个GPU的名称为/gpu:0,第二个GPU名称为/gpu:1,以此类推。

TensorFlow提供了一个快捷的方式来查看运行每一个运算的设备。在生成会话时,可以通过设置log_device_placement参数来打印运行每一个运算的设备。下面的程序展示了如何使用log_device_placement这个参数。

import tensorflow as tf

a = tf.constant([1.0, 2.0, 3.0], shape=[3], name='a')
b = tf.constant([1.0, 2.0, 3.0], shape=[3], name='b')
c = a + b
# 通过log_device_placement参数来输出运行每一个运算的设备。
sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
print sess.run(c)

'''
   在没有GPU的机器上运行以上代码可以得到以下输出:
Device mapping: no known devices.

add: /job:localhost/replica:0/task:0/cpu:0
b: /job:localhost/replica:0/task:0/cpu:0
a: /job:localhost/replica:0/task:0/cpu:0
[ 2.  4.  6.]
'''

在以上代码中,TensorFlow程序生成会话时加入了参数log_device_placement=True,所以程序会将运行每一个操作的设备输出到屏幕。于是除了可以看到最后的计算结果之外,还可以看到类似“add:/job:localhost/replica:0/task:0/cpu:0”这样的输出。这些输出显示了执行每一个运算的设备。比如加法操作add是通过CPU来运行的,因为它的设备名称中包含了/cpu:0。

在配置好GPU环境的TensorFlow中,如果操作没有明确地指定运行设备,那么TensorFlow会优先选择GPU。比如将以上代码在亚马逊(Amazon Web Services, AWS)的 g2.8xlarge实例上运行时,会得到以下运行结果。

Device mapping:
/job:localhost/replica:0/task:0/gpu:0 -> device: 0, name: GRID K520, pci bus id: 0000:00:03.0
/job:localhost/replica:0/task:0/gpu:1 -> device: 1, name: GRID K520, pci bus id: 0000:00:04.0
/job:localhost/replica:0/task:0/gpu:2 -> device: 2, name: GRID K520, pci bus id: 0000:00:05.0
/job:localhost/replica:0/task:0/gpu:3 -> device: 3, name: GRID K520, pci bus id: 0000:00:06.0

add: /job:localhost/replica:0/task:0/gpu:0
b: /job:localhost/replica:0/task:0/gpu:0
a: /job:localhost/replica:0/task:0/gpu:0
[ 2.  4.  6.]

从上面的输出可以看到在配置好GPU环境的TensorFlow中,TensorFlow会自动优先将运算放置在GPU上。不过,尽管g2.8xlarge实例有4个GPU,在默认情况下,TensorFlow只会将运算优先放到/gpu:0上。于是可以看见在上面的程序中,所有的运算都被放在了/gpu:0上。如果需要将某些运算放到不同的GPU或者CPU上,就需要通过tf.device来手工指定。下面的程序给出了一个通过tf.device手工指定运行设备的样例。

import tensorflow as tf

# 通过tf.device将运算指定到特定的设备上。
with tf.device('/cpu:0'):
   a = tf.constant([1.0, 2.0, 3.0], shape=[3], name='a')
   b = tf.constant([1.0, 2.0, 3.0], shape=[3], name='b')
with tf.device('/gpu:1'):
    c = a + b

sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
print sess.run(c)

'''
在AWS g2.8xlarge实例上运行上述代码可以得到一下结果:
Device mapping:
/job:localhost/replica:0/task:0/gpu:0 -> device: 0, name: GRID K520, pci bus id: 0000:00:03.0
/job:localhost/replica:0/task:0/gpu:1 -> device: 1, name: GRID K520, pci bus id: 0000:00:04.0
/job:localhost/replica:0/task:0/gpu:2 -> device: 2, name: GRID K520, pci bus id: 0000:00:05.0
/job:localhost/replica:0/task:0/gpu:3 -> device: 3, name: GRID K520, pci bus id: 0000:00:06.0

add: /job:localhost/replica:0/task:0/gpu:1
b: /job:localhost/replica:0/task:0/cpu:0
a: /job:localhost/replica:0/task:0/cpu:0
[ 2.  4.  6.]
'''

在以上代码中可以看到生成常量a和b的操作被加载到了CPU上,而加法操作被放到了第二个GPU“/gpu:1”上。在TensorFlow中,不是所有的操作都可以被放在GPU上,如果强行将无法放在GPU上的操作指定到GPU上,那么程序将会报错。

地址:http://www.infoq.com/cn/articles/introduction-of-tensorflow-part07?utm_source=articles_about_TensorFlow&utm_medium=link&utm_campaign=TensorFlow

您可能感兴趣的与本文相关的镜像

TensorFlow-v2.15

TensorFlow-v2.15

TensorFlow

TensorFlow 是由Google Brain 团队开发的开源机器学习框架,广泛应用于深度学习研究和生产环境。 它提供了一个灵活的平台,用于构建和训练各种机器学习模型

TensorFlow.NET 缺少 `tf.distribute.MirroredStrategy`,可考虑以下解决方法: ### 1. 检查版本兼容性 要保证使用的是支持多 GPU 分布式训练功能的 TensorFlow.NET 版本。可通过更新到最新版本来尝试解决,使用 NuGet 包管理器更新 TensorFlow.NET 到最新稳定版。在 Visual Studio 中,可在“解决方案资源管理器”里右键单击项目,选择“管理 NuGet 包”,搜索 `TensorFlow.NET` 并更新到最新版本。 ### 2. 实现自定义多 GPU 训练逻辑 若当前版本确实不支持 `tf.distribute.MirroredStrategy`,可自定义多 GPU 训练逻辑。以下是一个简单示例: ```csharp using Tensorflow; using Tensorflow.Keras; using static Tensorflow.Binding; // 假设已有模型定义 var model = new Sequential(); model.add(keras.layers.Dense(64, activation: "relu", input_shape: new TensorShape(784))); model.add(keras.layers.Dense(10, activation: "softmax")); // 编译模型 model.compile(optimizer: keras.optimizers.Adam(), loss: keras.losses.SparseCategoricalCrossentropy(), metrics: new[] { "accuracy" }); // 模拟多个 GPU var gpus = tf.config.list_physical_devices("GPU"); if (gpus.Length > 1) { // 这里可以实现自定义的多 GPU 数据分发和训练逻辑 // 例如将数据分割到不同 GPU 上进行计算 // 以下只是简单示例,实际需要更复杂的实现 foreach (var gpu in gpus) { using (tf.device(gpu.name)) { // 在该 GPU 上进行前向传播和反向传播等操作 // 这里简化为调用模型的 fit 方法 model.fit(x_train, y_train, epochs: 10, batch_size: 32); } } } else { model.fit(x_train, y_train, epochs: 10, batch_size: 32); } ``` ### 3. 社区支持与开源贡献 查看 TensorFlow.NET 的开源仓库,在 Issues 中搜索是否有关于 `tf.distribute.MirroredStrategy` 缺失的讨论。若没有,可创建一个新的 Issue 反馈该问题。也可考虑为项目做出开源贡献,尝试自己实现 `tf.distribute.MirroredStrategy` 并提交 Pull Request。
评论 2
成就一亿技术人!
拼手气红包6.0元
还能输入1000个字符
 
红包 添加红包
表情包 插入表情
 条评论被折叠 查看
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值