TensorFlow 指定 GPU 与动态分配显存设置

沈亦白

发布时间：2026-10-13 | 2人浏览过

原创

本文介绍如何在 TensorFlow 中指定使用的 GPU 并动态分配显存。通过命令行设置 CUDA_VISIBLE_DEVICES 环境变量可限制可见 GPU，使用 nvidia-smi 查看状态。在 Python 代码中，通过修改 os.environ 和配置 tf.ConfigProto 的 allow_growth 属性，实现按需分配显存，避免内存溢出。

通过命令行指定 GPU

在 TensorFlow 中，默认情况下会尝试占用所有可用的 GPU 资源。如果希望指定特定的 GPU 供程序使用，可以通过命令行设置环境变量来实现。


export CUDA_VISIBLE_DEVICES=1

例如，通过设置环境变量 CUDA_VISIBLE_DEVICES=1，可以限制程序仅使用编号为 1 的 GPU。执行该命令后，程序将只能看到并访问指定的 GPU 设备。


nvidia-smi

为了确认当前系统中各个 GPU 的使用状态和显存占用情况，可以使用 nvidia-smi 命令进行查看。该命令能显示 GPU 的实时负载、显存使用量以及正在运行的进程信息。

在 Python 代码中动态分配显存

此外，也可以在 Python 程序中直接指定使用的 GPU，并配置动态分配显存的功能。以下代码展示了如何在程序内部实现这一设置，确保 TensorFlow 仅按需分配显存，避免一次性占用过多资源。


import os
import sys
os.environ['CUDA_VISIBLE_DEVICES'] = sys.argv[1]
import tensorflow as tf
from keras.backend.tensorflow_backend import set_session
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
set_session(tf.Session(config=config))

通过上述代码，程序将使用 sys.argv[1] 指定的 GPU 设备，并启用动态显存分配功能。这意味着 TensorFlow 将根据实际计算需求逐步申请显存，而不是预先分配全部可用显存，从而更灵活地管理资源。

以上关于 TensorFlow 指定 GPU 与动态分配显存的设置方法，旨在帮助开发者更高效地利用 GPU 资源。希望这些内容能为您提供参考，如有更多问题，欢迎继续交流。