tensorflow jupyter 显存占用 释放 限制显存使用 多进程

jupyter 释放显存

jupyter 点击重启按钮会释放显存，如果没有，可以重启jupyter或者docker容器

设置显存动态增长

tensorflow 官方文档: https://www.tensorflow.org/guide/gpu#limiting_gpu_memory_growth

import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
assert len(gpus) > 0
print(gpus)
try:
    # Currently, memory growth needs to be the same across GPUs
    for i in range(len(gpus)):
        tf.config.experimental.set_memory_growth(gpus[i], True)
    logical_gpus = tf.config.experimental.list_logical_devices('GPU')
    print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
except Exception as e:
    # Memory growth must be set before GPUs have been initialized
    print("[error]", e)

查看显卡信息

用 pynvml模块来获取nvidia显卡信息, 通过pip install nvidia-ml-py3来安装

pynvml.nvmlInit()
gpu_num = pynvml.nvmlDeviceGetCount()
for i in range(gpu_num):
    h = pynvml.nvmlDeviceGetHandleByIndex(i)
    name = pynvml.nvmlDeviceGetName(h)
    info = pynvml.nvmlDeviceGetMemoryInfo(h)
    msg = "GPU:{}, used:{}/{}MB, free:{}MB".format(name, info.used/1024/1024, info.total/1024/1024, info.free/1024/1024)
    print(msg)
pynvml.nvmlShutdown()

TensorFlow 显存释放

如果要运行完成后释放显卡内存, 结束进程即可,会自动释放

• 对于一个独立的程序,运行完成退出即被释放了
• 对于程序里面调用训练,训练完成后主程序不退出的情况, 需要在一个新的子进程中运行训练,练完后结束进程即可. 注意主程序退出要把子进程给退出或者杀死,以免出现孤儿进程. 另外需要注意， 引入 tensorflow 时不要直接在全局 import,而是在需要的时候，或者在新的进程入口处引入，而且父进程不要import, 父进程只做管理，不然由于在父进程中初始化了cuda，在创建子进程时会使用内存拷贝的信息，导致cuda在子进程中初始化失败（could not retrieve CUDA device count: CUDA_ERROR_NOT_INITIALIZED）。如果父进程一定要调用tensorflow，比如父进程不需要训练但是需要通过tensorflow库获取显卡信息，可以开多个子进程来解决这个问题

参考

• tensorflow 官方文档: https://www.tensorflow.org/guide/gpu#limiting_gpu_memory_growth