TensorFlow禁用GPU和动态分配显存

天边一坨浮云

已于 2022-08-22 10:58:38 修改

阅读量935

点赞数

分类专栏： # TensorFlow 文章标签： tensorflow

于 2022-06-03 09:49:02 首次发布

本文链接：https://blog.csdn.net/ytomc/article/details/125110190

版权

TensorFlow 专栏收录该内容

13 篇文章 5 订阅

订阅专栏

禁用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