TensorFlow禁用GPU和动态分配显存

禁用GPU

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

动态分配显存

第一种方式大家应该都知道，主要是给Session用的

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

但我使用的是Tensorflow的Estimator，可以采用以下方法

session_config = tf.ConfigProto(log_device_placement=True,allow_soft_placement=True)
session_config.gpu_options.per_process_gpu_memory_fraction = 0.8
session_config.gpu_options.allow_growth = True # 自适应
run_config = tf.estimator.RunConfig(
        session_config=session_config,
        model_dir=FLAGS.output_dir,
        save_checkpoints_steps=FLAGS.save_checkpoints_steps)

如果是TPU，也是一样，无非就是tf.estimator改为tf.estimator.tpu