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数目