python 设置使用独立显卡_限制TensorFlow显存使用量-CSDN博客

问题是这样的：比如我使用一块英伟达的P40，每张卡都是24GB的显存，但是我的程序只需要使用4G，深度学习框架是TensorFlow，那么我怎么限制这个显存的使用量呢？

一、手动设置使用比例

众所周知，TensorFlow是不会设定显存大小的，比如程序4G显存也能跑起来，但是TensorFlow就是会把P40的24GB全部用光，这显然是很蛋疼的，所以一般会通过手动设置GPU使用比例，参考https://blog.csdn.net/m0_37041325/article/details/77488981，可以这样设置：

gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.17)
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))

这个0.17就是我手动设置的4/24的大小，这样设置好就能够防止把资源吃光了（土豪随意）。

二、自动设置使用比例

但这样做有个显然的缺点，就是这个比例的设置严重依赖显卡的大小，因为程序所需的显存是不变的，即分子不变，那么使用比例就跟分母即显存大小相关了。

设想这样一个场景：我本地有张1060的卡，只有6GB，那么可以设置一个4/6=0.67的使用比例，本地我只想验证一下程序能不能跑通，跑通了我就放到P40的服务器了，这样又要设置一个4/24=0.17的使用比例了，这样就有两个使用比例了，还是同一个程序。

一个简单的方法就是读取使用平台的名称，即platform.node()，根据这个来判断是哪个机器，就使用对应的GPU使用比例了。

但这个显然不够smart了，而且，当面临第三个平台，这种不智能化的部署会出问题的，我想到的解决方法就是先查看显卡的大小，然而根据这个大小来确定使用比例。

查看显卡大小有两种方法：

第一种可以参考使用Tensorflow或Keras时对GPU内存限制 | 沐雨浥尘及python实现自动查找并使用空闲GPU，直接使用官方命令取得，好处是不用安装新的库：

# compute GPU memory to use
memoryList = list(map(int, os.popen("nvidia-smi -q -d Memory | grep -A4 GPU | grep Total | awk '{print $3}'").readlines()))
memoryTotal = memoryList[0]
print ("memoryTotal = %s" % memoryTotal)	# 22912
memoryLimited = 4000.0	# use 4G memory
configTensorFlow.gpu_options.per_process_gpu_memory_fraction = memoryLimited / memoryTotal

第二种是python查看显卡gpu信息，貌似要安装一个Python库，没有实践过，就不多说了，有兴趣的可以实践实践。

【已完结】