- 作者:
- 分类:知识&开发->AI->工具
- 阅读:3345
- 点赞:8
- 版权:CC BY-SA 4.0
- 创建:2020-03-25
- 更新:2020-10-19
Tensorflow 限制显存使用和释放显存的方法,以及多进程的使用
版权声明:本文为 neucrack 的原创文章,遵循 CC 4.0 BY-SA 版权协议,转载请附上原文出处链接及本声明。
原文链接(持续更新):https://neucrack.com/p/257
原文链接(持续更新):https://neucrack.com/p/257
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