TensorFlow指定CPU与GPU运算的方法实现

追风的人
追风的人
发布时间:2026-10-13| 5人浏览过
原创

本文介绍如何在TensorFlow中指定CPU与GPU运算。通过device语句明确指派设备,利用ConfigProto配置日志与软放置。详解GPU资源分配策略,包括按需增长(allow_growth)与固定比例(per_process_gpu_memory_fraction),确保显存高效利用。

指定GPU运算

1.指定GPU运算

若安装的是GPU版本,运行过程中TensorFlow会自动检测。一旦检测到GPU,TensorFlow会尽可能利用找到的第一个GPU来执行操作。

如果机器上存在多个可用GPU,除第一个外的其他GPU默认不参与计算。为使TensorFlow使用这些GPU,必须将OP明确指派给它们执行。使用with......device语句可以指派特定的CPU或者GPU执行操作:

设备指派示例


import tensorflow as tf
import numpy as np

with tf.Session() as sess:
  with tf.device('/cpu:0'):
    a = tf.placeholder(tf.int32)
    b = tf.placeholder(tf.int32)
    add = tf.add(a, b)
    sum = sess.run(add, feed_dict={a: 3, b: 4})
    print(sum)

设备的字符串标识,当前支持的设备包括以下的几种:

cpu:0 机器的第一个cpu。

gpu:0 机器的第一个gpu,如果有的话

gpu:1 机器的第二个gpu,依次类推

类似的还有tf.ConfigProto来构建一个config,在config中指定相关的GPU,并且在session中传入参数config=“自己创建的config”来指定gpu操作

其中,tf.ConfigProto函数的参数如下:

log_device_placement=True: 是否打印设备分配日志

allow_soft_placement=True: 如果指定的设备不存在,允许TF自动分配设备

配置ConfigProto


import tensorflow as tf
import numpy as np

config = tf.ConfigProto(log_device_placement=True, allow_soft_placement=True)

with tf.Session(config=config) as sess:
  a = tf.placeholder(tf.int32)
  b = tf.placeholder(tf.int32)
  add = tf.add(a, b)
  sum = sess.run(add, feed_dict={a: 3, b: 4})
  print(sum)

设置GPU使用资源

2.设置GPU使用资源

上文的tf.ConfigProto函数生成的config之后,还可以设置其属性来分配GPU的运算资源,如下代码就是按需分配

按需分配显存


import tensorflow as tf
import numpy as np

config = tf.ConfigProto(log_device_placement=True, allow_soft_placement=True)
config.gpu_options.allow_growth = True

with tf.Session(config=config) as sess:
  a = tf.placeholder(tf.int32)
  b = tf.placeholder(tf.int32)
  add = tf.add(a, b)
  sum = sess.run(add, feed_dict={a: 3, b: 4})
  print(sum)

使用allow_growth option,刚开始会分配少量的GPU容量,然后按需要慢慢的增加,有与不会释放内存,随意会导致内存碎片。

同样,上述的代码也可以在config创建时指定,

创建时指定配置


import tensorflow as tf
import numpy as np

gpu_options = tf.GPUOptions(allow_growth=True)
config = tf.ConfigProto(gpu_options=gpu_options)


with tf.Session(config=config) as sess:
  a = tf.placeholder(tf.int32)
  b = tf.placeholder(tf.int32)
  add = tf.add(a, b)
  sum = sess.run(add, feed_dict={a: 3, b: 4})
  print(sum)

我们还可以给gpu分配固定大小的计算资源。

固定大小资源分配


gpu_options = tf.GPUOptions(allow_growth=True, per_process_gpu_memory_fraction=0.5)

上述代码的含义是分配给tensorflow的GPU显存大小为:GPU的实际显存*0.5

到此这篇关于tensorflow指定CPU与GPU运算的方法实现的文章就介绍到这了,更多相关tensorflow指定CPU与GPU运算内容请搜索本站以前的文章或继续浏览下面的相关文章希望大家以后多多支持本站!

相关标签:

免责声明:本文内容仅供参考,不代表本站观点。转载或引用请注明出处。

Python 相关产品和应用

全部下载
Selenium

Selenium

4.51.0 · Windows

用于浏览器自动化和 Web 测试的 Python 工具。

scikit-learn

scikit-learn

1.9.1 · Windows

常用 Python 机器学习算法与模型评估工具库。

Black

Black

26.10.1 · Windows

强调一致性的 Python 自动代码格式化工具。

Ruff

Ruff

0.17.1 · macOS

由 Rust 编写的高速 Python 代码检查和格式化工具。

pytest

pytest

9.1.1 · Windows

简洁而强大的 Python 自动化测试框架。

FastAPI

FastAPI

0.146.0 · Windows

高性能 Python API 框架,支持类型提示和自动接口文档。

Flask

Flask

3.1.3 · Windows

轻量灵活的 Python Web 应用框架。

Django

Django

6.1.2 · Windows

成熟的 Python Web 框架,适合构建内容站点和业务系统。

Spyder

Spyder

6.1.7 · Windows

面向科学计算与数据分析的 Python 集成开发环境。

SciPy

SciPy

1.18.1 · Windows

面向科学、工程和技术计算的 Python 算法库。