TensorFlow使用显存的设置

本文介绍了如何在TensorFlow中管理显存使用,包括设置显存使用比例、按需分配显存以及限制可见GPU数量的方法。

TensorFlow默认情况下训练会占用全部显存。
可以设置参数来修改显存使用情况。

设置使用显存比例

在构造tf.Session()时候通过传递tf.GPUOptions作为可选配置参数的一部分来显式地指定需要分配的显存比例,如下所示:

假如有12GB的显存并使用其中的4GB:

gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)  
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))  

per_process_gpu_memory_fraction指定了每个GPU进程中使用显存的上限,但它只能均匀作用于所有GPU,无法对不同GPU设置不同的上限.

设置显存使用随需求增加

尝试如下设置:

config = tf.ConfigProto()  
config.gpu_options.allow_growth=True  
sess = tf.Session(config=config) 

allow_growth设置为True时,分配器将不会指定所有的GPU内存,而是根据需求增长

设置GPU的数目

在执行训练脚本前使用:

export CUDA_VISIBLE_DEVICES=1  

来限制可见的GPU数目

### TensorFlow 配置 GPU 显存使用 为了优化 TensorFlow 的性能并充分利用可用的 GPU 资源,可以通过调整 `tf.config` 来控制 GPU 显存分配策略。以下是几种常见的方法来增加 GPU 存使用的配置方式。 #### 方法一:设置动态显存增长 通过启用动态显存增长功能,可以让 TensorFlow 按需申请显存而不是一次性占用全部显存。这有助于提高多任务环境下的资源利用率[^3]。 ```python import tensorflow as tf gpus = tf.config.experimental.list_physical_devices('GPU') if gpus: try: for gpu in gpus: tf.config.experimental.set_memory_growth(gpu, True) except RuntimeError as e: print(e) ``` 上述代码会遍历所有检测到的物理 GPU 设备,并为每个设备开启动态显存增长模式。 #### 方法二:手动指定显存量 如果希望精确控制每块 GPU 使用的具体显存量,则可以采用虚拟设备的方式进行划分。这种方法适用于需要严格限制单个进程所消耗的最大显存量的情况[^4]。 ```python gpu_options = [ {'limit': '4GB'}, # 设置第一个GPU最大可使用4G显存 ] for i, option in enumerate(gpu_options): with tf.device(f'/device:GPU:{i}'): logical_gpus = tf.config.experimental.get_visible_devices('GPU', index=i) if not logical_gpus: continue vram_limit = int(option['limit'].replace('GB','')) * (1 << 30) tf.config.experimental.set_virtual_device_configuration( logical_gpus[0], [tf.config.LogicalDeviceConfiguration(memory_limit=vram_limit)] ) ``` 此脚本定义了一个列表用于存储各张卡上的VRAM上限值(单位转换成字节),随后逐一应用这些设定给对应的逻辑设备上。 #### 方法三:完全禁用其他计算引擎 当运行环境中存在多个框架共存时(比如PyTorch),可能会导致不必要的竞争从而降低整体效率。因此,在某些情况下可能还需要考虑关闭其它库对于相同硬件的支持访问权限[^5]。 例如,在程序开头加入如下命令即可阻止 Pytorch 对 NVIDIA CUDA 的调用: ```bash export LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libcuda.so.1 ``` 或者直接修改 Python 启动参数: ```bash CUDA_VISIBLE_DEVICES="" python your_script.py ``` 以上操作均能有效减少潜在干扰因素的影响程度。 --- ### 总结 综上所述,针对不同场景需求提供了三种主要途径去提升 Tensorflow 下面 GPU Memory Usage 效率问题解决方案;分别是启用了自动扩展机制、自定义化分区管理以及屏蔽外部干扰源三个方面进行了深入探讨分析说明[^6]。
评论 1
成就一亿技术人!
拼手气红包6.0元
还能输入1000个字符
 
红包 添加红包
表情包 插入表情
 条评论被折叠 查看
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

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

抵扣说明:

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

余额充值