TensorFlow 管理GPU

项目开发中，发现pytorch对显存管理真的是很差，一不小心就会OOM了，感觉好坑，而且会强行分配，造成显存卡住，python处于假死状态，即使你有supervisor监控也没有用。（当然这种情况需要基于内容的监控来解决假死问题，Ping Test + XPing Test）

TensorFlow对显存的管理相对于pytorch要好很多，支持几种不同的显存管理方法。

1. 我用的NVIDIA的GPU，想查看GPU的情况，可以使用命令

nvidia-smi

如果想定时刷新GPU状态，可以用命令

nvidia-smi -l

但是命令“nvidia-smi -l”在同一个屏幕上平铺式刷新，不太美观，可以用命令0.1秒刷新一次（你可以改刷新频率）

watch -n 0.1 nvidia-smi

2. TensorFlow对显存管理

2.1 启动程序时在terminal里指定程序使用显卡号

CUDA_VISIBLE_DEVICES=<device_id> python <your_python_file.py>

例如要指定卡0来跑你的test.py程序，可以使用如下命令：

CUDA_VISIBLE_DEVICES=0 python test.py

如果要指定多张卡，可以把device_id用逗号“,”分割，例如要0和1号卡跑test.py程序

CUDA_VISIBLE_DEVICES=0,1 python test.py

当然你也可以把卡号用双引号括起来，但是双引号不是必须的

CUDA_VISIBLE_DEVICES="0,1" python test.py

2.2 在python代码中指定GPU

你需要引入os这个package，然后在你使用GPU之前执行这个就可以，一般放程序的开头

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"

2.3 给每个进程分配固定的显存

这里你程序启动的时候时候要加载到GPU上的时候，预先分配设置的大小，如果不够，那就直接OOM了。所以如果跑服务的话，最好不要这么搞；如果做训练，可以通过这种方式强制GPU memory。

config = tf.ConfigProto() 
# set gpu memory as 30% for each process, you can startup multiple-process for your program
config.gpu_options.per_process_gpu_memory_fraction = 0.3 
session = tf.Session(config=config)

2.4 给每个进程按需自动增加

当设置gpu_options.allow_growth为True时，程序会按需去申请GPU memory，而不至于独占显存，这个与预先分配显存的方式各有用处，根据自己的需要设定就好了。

config = tf.ConfigProto() 
config.gpu_options.allow_growth = True 
session = tf.Session(config=config)

可以参考实际例子mnist_cnn对GPU的使用，你可以尝试四种情况的GPU使用情况。