tensorflow报错 或者 keras报错 以及tf.keras报错： OOM 显存不足

Posted 2021-12-07 _刘文凯_

tensorflow报错 或者 keras报错，以及tf.keras报错：Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info.

文章目录

• • • 问题描述
    • 原因
    • 解决
    • • • • （1）在tensorflow中
          • （2）在keras中
          • （3）在tf.keras中

问题描述

今天进行5折交叉验证 以及 格点搜索时遇到的问题，以前数据量过大、或者bitch_size过大时也会发生，具体如图所示：
在linux中使用命令：watch -n 0.1 nvidia-smi 查看GPU使用情况

原因

由于显存不足造成的，但是并非真正的显存不足，而是由于tensorflow把显存吃完了，但是并没有实际的有效利用。因此可以分配需要的显存给tensorflow。（keras基于tensorflow的，同样适用）

解决

1、设置小的bitch_size， 虽然可以用，但是指标不治本。
2、手动设置GPU,在train.py中：

（1）在tensorflow中

import tensorflow as tf
import os

os.environ["CUDA_VISIBLE_DEVICES"] = "0" 指定使用哪个GPU
config = tf.ConfigProto()
config.gpu_options.allow_growth = True  # 按需分配显存
config.gpu_options.per_process_gpu_memory_fraction = 0.4  # 最大使用显存40%
session = tf.Session(config=config))  # 创建tensorflow会话
...

（2）在keras中

import tensorflow as tf
from keras.models import Sequential
import os
from keras.backend.tensorflow_backend import set_session ## 和tf.keras不同

os.environ["CUDA_VISIBLE_DEVICES"] = "0"
config = tf.ConfigProto()
config.gpu_options.allow_growth = True  # 按需分配显存
set_session(tf.Session(config=config)) # 把设置传给keras

model = Sequential()
...

（3）在tf.keras中

import tensorflow as tf
from tensorflow.keras.models import Sequential

import os
from tensorflow_core.python.keras.backend import set_session # 和keras不同

os.environ["CUDA_VISIBLE_DEVICES"] = "0"
config = tf.ConfigProto()
config.gpu_options.allow_growth = True  # 按需分配显存
config.gpu_options.per_process_gpu_memory_fraction = 0.4  # 最大使用显存40%
set_session(tf.Session(config=config))  # 把设置传给 tf.keras

model = Sequential()
...

补充：
tf.keras可以使用数据读取多线程加速：

model.fit(x_train，y_train，use_multiprocessing=True, workers=4) # 开启多线程，使用4个CPU

清空会话:

from tensorflow import keras
keras.backend.clear_session() 

清空后可以继续创建新的会话