在tensorflow中,tf.config.experimental中的方法很有可能会随着版本的提升改动或删除
目录
1 检查GPU是否可用
2 获取当前程序可使用的GPU列表
3 获取当前程序可使用的CPU列表
4 设置程序可见的范围
5 设置显存使用策略
5.1 程序用多少给多少,多的不给
5.2 给固定大小的显存
1 检查GPU是否可用
如果可用会返回True,不可以则返回False
2 获取当前程序可使用的GPU列表
也有一种情况是机器上有GPU,但是在这里不可见,这个时候就要检查tensorflow版本,比如说如果装的是CPU版本就不可见,之后还有显卡驱动这些也要看一看
我当前只有一个GPU,所以列表中只有一个元素
3 获取当前程序可使用的CPU列表
我当前只有一个CPU,所以列表中只有一个元素
4 设置程序可见的范围
GPU只要让程序看见,程序就会默认使用所有GPU
我们上面获取了GPU的列表,如果我们有多块GPU的话,并且我们想使用多块GPU,那么我们就使用下面这个方法,比如说我们有两块,我现在要用这两块
如果是有多块GPU,但是我不想让程序使用,那这个时候就对着列表看,改一下上面[0:2]这个部分,如果有两块想用一块,那就改成[0:1]或[1:2],想用前面那块就[0:1],用后面那块就[1:2]
我们也可以通过控制环境变量来控制程序的可见范围,比如我现在有四块显卡,我想都投入使用
如果我们想少用哪一个,我们去掉对应的数字就可以
5 设置显存使用策略
在默认情况下,程序会占用所有可见GPU的所有显存
5.1 程序用多少给多少,多的不给
我们使用set_memory_growth(),让程序在可见gpu的范围内随用随取
5.2 给固定大小的显存
下面涉及到显存的大小了,我们先说一下显存在哪看,我们可以使用nvidia-smi来查看设备的显存
我当前的显卡为 GeForce GTX 970M ,显存为3072MiB,此时使用的显存为39MiB,如果我们的设备上有多个显卡,在这里会显示多行数据
一般来讲我们对一块进行限制,我们当前想仅让显卡使用1024MiB(1GiB)大小的显存,那么我们这样写
如果我现在有两块显卡,我现在都限制其只使用1024MiB,那么我们这样写