TensorFlow

TensorFlow与显存

开始学习用TensorFlow,这东西与Theano不同,默认情况下, 在开启一个Session后,几乎占用显卡的所有显存。如果同一个机器、显卡多个人使用,基本上就是先到先得,后来的程序会崩溃。查了下文档有两种方法控制显存:

第一种是设置成预加载比例:

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

还有一种是自适应,需要多少就占多少:

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