【使用 jupyter notebook 进行网络训练时显存不足,设置GPU占用率】

使用 Jupyter Notebook 进行长期模型训练时,可能由于 GPU 内存无法及时释放导致报错,可以调小batch_szie并且及时释放jupyter notebook占用的内存。

1 用到的方法和对象

config=tf.compat.v1.ConfigProto()  #返回一个session实例对象

config.gpu_options.allow_growth = True  # 配置GPU内存分配方式,按需增长
config.gpu_options.per_process_gpu_memory_fraction = 0.6  #限制GPU占用率 

set_session(tf.compat.v1.Session(config=config))#将上面的配置进行应用,设置全局会话

2 具体实现如下:只需在训练开始前调用该函数即可,其中GPU占用率可自行设置

from keras.backend import set_session
from keras.backend import clear_session
from keras.backend import get_session
import tensorflow as tf
import gc #此模块提供可选的垃圾回收器的接口
 #使用 Jupyter Notebook 进行长期模型训练时,可能由于 GPU 内存无法及时释放导致报错
#会话(Session)可以管理*TensorFlow运行时的所有资源
def set_gpu():
    sess = get_session() #返回后端使用的TF会话。
    clear_session()#释放全局状态即重置keras生成的所有状态:这有助于避免旧模型和层的混乱,特别是在内存有限的时候。
    sess.close()#关闭会话
 
    try:
        del classifier 
    except:
        pass
 
    print(gc.collect()) #若被调用时不包含参数,则启动完全的垃圾回收 
 
    config = tf.compat.v1.ConfigProto()#返回一个session实例对象
    # 配置GPU内存分配方式,按需增长,很关键
    config.gpu_options.allow_growth = True
    config.gpu_options.per_process_gpu_memory_fraction = 0.6#限制GPU占用率 
    ''' GPU可见设备到虚拟设备的映射,为逗号分割的列表,
    本机GPU设备0映射到虚拟device:GPU:0 '''
    config.gpu_options.visible_device_list = "0" 
    set_session(tf.compat.v1.Session(config=config))#将上面的配置进行应用,设置全局会话

你可能感兴趣的:(jupyter,python,ide)