TensorFlow 显存管理：按比例分配与自适应增长配置指南

星野

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

原创

本文详解在 TensorFlow 中通过 ConfigProto 配置 GPU 显存占用策略。涵盖按比例固定分配（per_process_gpu_memory_fraction）与自适应增长（allow_growth）两种模式，并介绍如何利用 CUDA_VISIBLE_DEVICES 环境变量在多台 GPU 环境中指定计算设备，帮助优化内存使用并避免 OOM 错误。

TensorFlow GPU 显存配置概述

在 TensorFlow 框架下运行模型时，合理设置显存占用策略至关重要。

按比例分配显存

1. 按比例分配显存


config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.4 # 根据自己的需求确定
session = tf.Session(config=config, ...)

启用显存自适应增长

2. 自适应增长模式


config = tf.ConfigProto()
config.gpu_options.allow_growth = True
session = tf.Session(config=config, ...)

ConfigProto 配置方法

配置 GPU 使用率时，需在创建 Session 阶段对 config 类进行相应设置。

多 GPU 环境下的设备指定

此外，当计算机配备多块 GPU 时，可以通过指定选取哪一块 GPU 进行计算来优化资源分配。


# 在程序开头添加
os.environ['CUDA_VISIBLE_DEVICES'] = '0' #使用 GPU 0 0对应着ubuntu系统给GPU的序号,可通过Nvidia-smi命令查看

若系统中存在多个 GPU，可通过环境变量限制可见设备。


os.environ['CUDA_VISIBLE_DEVICES'] = '0,1' # 使用 GPU 0，1

以下是一个在代码中指定 GPU 使用的常见范例：


import tensorflow as tf
import numpy as np
import os

os.environ['CUDA_VISIBLE_DEVICES']='0'
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction=0.6

with tf.Session(graph=...,config=config) as sess:

## 后续的操作

以上关于 TensorFlow 设置显存自适应及显存比例的操作即为全部内容，希望能为您提供参考。