选中内容(绿色)时除了会搜索文章名,还会搜索文章内容
点击结果中的文章名进入文章界面后可以按Ctrl+F在页面内搜索
  • 版权:CC BY-SA 4.0
  • 创建:2020-03-25
  • 更新:2020-10-19
Tensorflow 限制显存使用和释放显存的方法,以及多进程的使用


jupyter 释放显存

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

设置显存动态增长

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

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

查看显卡信息

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

  1. pynvml.nvmlInit()
  2. gpu_num = pynvml.nvmlDeviceGetCount()
  3. for i in range(gpu_num):
  4. h = pynvml.nvmlDeviceGetHandleByIndex(i)
  5. name = pynvml.nvmlDeviceGetName(h)
  6. info = pynvml.nvmlDeviceGetMemoryInfo(h)
  7. msg = "GPU:{}, used:{}/{}MB, free:{}MB".format(name, info.used/1024/1024, info.total/1024/1024, info.free/1024/1024)
  8. print(msg)
  9. pynvml.nvmlShutdown()

TensorFlow 显存释放

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

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

参考

文章有误?有想法想讨论?查看或者发起勘误/讨论 主题
(发起评论需要先登录 github)

/wallpaper/wallhaven-pk9kve.jpg