使用tensorflow,如果不加设置,即使是很小的模型也会占用整块GPU,造成资源浪费。
所以我们需要设置,使程序按需使用GPU。
具体设置方法:
gpu_options = tf.GPUOptions(allow_growth=True)
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))
# sess = tf.InteractiveSession(config=tf.ConfigProto(gpu_options=gpu_options))
gpu_no = '' # or '1'
os.environ["CUDA_VISIBLE_DEVICES"] = gpu_no
# 定义TensorFlow配置
config = tf.ConfigProto()
# 配置GPU内存分配方式,按需增长,很关键
config.gpu_options.allow_growth = True
# 配置可使用的显存比例
config.gpu_options.per_process_gpu_memory_fraction = 0.1
# 在创建session的时候把config作为参数传进去
sess = tf.InteractiveSession(config = config)
gpu_options = tf.GPUOptions(allow_growth=True)
# sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))
with tf.Session(config=tf.ConfigProto(gpu_options=gpu_options)) as sess:
说明:使用jupyter notebook,如果没有上述设置,整个GPU会一直被占用;使用pycharm,虽然GPU不会被一直占用,但是运行时还是会占用整个GPU,所以也需要上述设置。
手机扫一扫
移动阅读更方便
你可能感兴趣的文章