tensorflow显存设置

MagiChos

于 2017-12-28 10:16:36 发布

阅读量447

点赞数

CC 4.0 BY-SA版权

分类专栏： tensorflow

本文链接：https://blog.youkuaiyun.com/hanshuobest/article/details/78919023

tensorflow 专栏收录该内容

9 篇文章

订阅专栏

本文介绍了一种在TensorFlow中限制GPU显存占用的方法，通过设置程序最多使用指定GPU的90%显存，并允许程序按需申请内存，从而有效避免了训练过程中显存被完全占用的问题。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

为避免训练时占用全部的gpu显存，只需设置如下代码

import numpy as np
import tensorflow as tf
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
# 程序最多占用制定gpu的50%显存
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.9)
# 程序按需申请内存
gpu_options.allow_growth = True
config = tf.ConfigProto(gpu_options=gpu_options)

init = tf.global_variables_initializer()
sess = tf.Session(config=config)

sess.close()