TensorFlow 训练网络时 指定gpu的使用率

版权声明：聂CC~ https://blog.csdn.net/ncc1995/article/details/88072516

训练网络时，TensorFlow会占满gpu显存，在我的电脑上就是网络跑不起来，内存不够用，今天又遇到了这个情况，所以总结一下解决这个问题的方法。

用TensorFlow训练网络时对gpu的设置：

1. 指定固定的gpu，我这里只有一个就不用进行这一步了。
2. 在每个gpu上，让gpu动态申请内存，用多少申请多少。
3. 直接限制gpu的使用率，自己设置数值

一、指定gpu

1、用环境变量指定

目前有0，1两个个GPU，找到空闲的GPU号，可以使用环境变量CUDA_VISIBLE_DEVICES：

环境变量的定义格式： 
CUDA_VISIBLE_DEVICES=1 
CUDA_VISIBLE_DEVICES=0,1 

运行时输入以下命令运行程序：使用空闲的GPU 0

$ CUDA_VISIBLE_DEVICES=0 python trian.py 

2、直接在代码中指定

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2"

二、动态申请内存

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

三、限定gpu使用率

gpu_options=tf.GPUOptions(per_process_gpu_memory_fraction=0.4)  
config=tf.ConfigProto(gpu_options=gpu_options)  
session = tf.Session(config=config)  

 

参考链接：

https://blog.csdn.net/A632189007/article/details/77978058

https://blog.csdn.net/m0_37987687/article/details/80240998