【TF】显存分配问题

如果使用TF跑程序遇到了“Blas GEMM launch failed“的问题，那么一定是显存不够的原因。主要原因是因为TF在跑GPU代码时会默认申请全部的GPU资源！

1. 指定某块GPU

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

这其实是CUDA本身的参数，对所有深度学习框架都是适用的。

需要注意的是最好写在improt keres 和 import tensorflow之前，不然可能出错。

注意：写在程序的开头

2. 指定多块GPU

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

注意：写在程序的开头

3. 不使用GPU（使用CPU）

import os
os.environ["CUDA_VISIBLE_DEVICES"]="-1"

注意：写在程序的开头

4. 控制GPU显存使用比例

4.1 TF1.X

tf_config = tensorflow.ConfigProto()  
tf_config.gpu_options.per_process_gpu_memory_fraction = 0.5 # 分配50%  
session = tensorflow.Session(config=tf_config)

注意：写在程序的开头

4.2 TF2.X

gpus = tf.config.experimental.list_physical_devices('GPU')
# 对需要进行限制的GPU进行设置
tf.config.experimental.set_virtual_device_configuration(gpus[0],
                                                        [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=8192)])

注意：写在程序的开头

5. 指定GPU+控制显存使用

5.1 TF1.x

tf_config = tensorflow.ConfigProto()  
tf_config.gpu_options.allow_growth = True  
session = tensorflow.Session(config=tf_config) 

5.2 TF2.x

config = tf.compat.v1.ConfigProto(gpu_options=tf.compat.v1.GPUOptions(allow_growth=True))
sss = tf.compat.v1.Session(config=config)

参考文章:

https://www.cnblogs.com/youyou0/p/9547751.html
https://zhuanlan.zhihu.com/p/65218239