TensorFlow指定GPU或CPU执行相应运算

最新推荐文章于 2025-07-05 21:11:59 发布

翻译最新推荐文章于 2025-07-05 21:11:59 发布 · 1.5k 阅读

TensorFlow基础专栏收录该内容

52 篇文章

订阅专栏

本文详细介绍了如何在TensorFlow中配置设备，包括使用CPU和GPU进行操作。通过具体实例展示了如何利用tf.device函数指定默认设备，适用于Windows10环境下Python3.6.3与TensorFlow1.12版本。

部署运行你感兴趣的模型镜像

一、环境

TensorFlow API r1.12

CUDA 9.2 V9.2.148

cudnn64_7.dll

Python 3.6.3

Windows 10

二、官方说明

https://www.tensorflow.org/api_docs/python/tf/device

使用默认图（default graph）是Graph.device()的包装器（wrapper）

详细参照tf.Graph.device

输入：

在上下文中使用的设备名称或者函数

输出：

为新创建的操作指定默认设备的上下文管理器

三：实例

（1）使用CPU

with tf.device("/cpu:0"):
    embedding = tf.get_variable("embedding", [vocab_size, size], dtype=tf.float32)
    inputs = tf.nn.embedding_lookup(embedding, input_.input_data)

（2）使用GPU

with tf.device("/gpu:0"):
    embedding = tf.get_variable("embedding", [vocab_size, size], dtype=tf.float32)
    inputs = tf.nn.embedding_lookup(embedding, input_.input_data)

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