tensorflow报错显存不够的应对措施

发表于 分类于 阅读次数：

简单粗暴

1
2
3
4
5
6
7
# tensorflow 内存不够的情况 适用于tensorflow 2.x
import tensorflow as tf

os.environ["CUDA_VISIBLE_DEVICES"] = "0"  # 指定显卡
config = tf.compat.v1.ConfigProto(allow_soft_placement=True)  # 开启软放置，OP放到CPU上了，为啥源码默认值是7？
config.gpu_options.per_process_gpu_memory_fraction = 0.3  # 进程最多采用30%显存，默认是1，完整使用
tf.compat.v1.keras.backend.set_session(tf.compat.v1.Session(config=config))