在tensorflow中,tf.config.experimental中的方法很有可能会随着版本的提升改动或删除

目录

1  检查GPU是否可用

2  获取当前程序可使用的GPU列表

3  获取当前程序可使用的CPU列表

4  设置程序可见的范围

5  设置显存使用策略

5.1  程序用多少给多少,多的不给

5.2  给固定大小的显存


1  检查GPU是否可用

sd 里 cuda内存分配器 pytorch 的区别 gpu内存分配_tensorflow

如果可用会返回True,不可以则返回False

sd 里 cuda内存分配器 pytorch 的区别 gpu内存分配_人工智能_02

2  获取当前程序可使用的GPU列表

也有一种情况是机器上有GPU,但是在这里不可见,这个时候就要检查tensorflow版本,比如说如果装的是CPU版本就不可见,之后还有显卡驱动这些也要看一看

sd 里 cuda内存分配器 pytorch 的区别 gpu内存分配_显卡驱动_03

我当前只有一个GPU,所以列表中只有一个元素

sd 里 cuda内存分配器 pytorch 的区别 gpu内存分配_tensorflow_04

3  获取当前程序可使用的CPU列表

sd 里 cuda内存分配器 pytorch 的区别 gpu内存分配_环境变量_05

我当前只有一个CPU,所以列表中只有一个元素

sd 里 cuda内存分配器 pytorch 的区别 gpu内存分配_环境变量_06

4  设置程序可见的范围

GPU只要让程序看见,程序就会默认使用所有GPU

我们上面获取了GPU的列表,如果我们有多块GPU的话,并且我们想使用多块GPU,那么我们就使用下面这个方法,比如说我们有两块,我现在要用这两块

sd 里 cuda内存分配器 pytorch 的区别 gpu内存分配_python_07

如果是有多块GPU,但是我不想让程序使用,那这个时候就对着列表看,改一下上面[0:2]这个部分,如果有两块想用一块,那就改成[0:1]或[1:2],想用前面那块就[0:1],用后面那块就[1:2]

我们也可以通过控制环境变量来控制程序的可见范围,比如我现在有四块显卡,我想都投入使用

sd 里 cuda内存分配器 pytorch 的区别 gpu内存分配_python_08

如果我们想少用哪一个,我们去掉对应的数字就可以

5  设置显存使用策略

在默认情况下,程序会占用所有可见GPU的所有显存

5.1  程序用多少给多少,多的不给

我们使用set_memory_growth(),让程序在可见gpu的范围内随用随取

sd 里 cuda内存分配器 pytorch 的区别 gpu内存分配_人工智能_09

5.2  给固定大小的显存

下面涉及到显存的大小了,我们先说一下显存在哪看,我们可以使用nvidia-smi来查看设备的显存

我当前的显卡为 GeForce GTX 970M ,显存为3072MiB,此时使用的显存为39MiB,如果我们的设备上有多个显卡,在这里会显示多行数据

sd 里 cuda内存分配器 pytorch 的区别 gpu内存分配_显卡驱动_10

一般来讲我们对一块进行限制,我们当前想仅让显卡使用1024MiB(1GiB)大小的显存,那么我们这样写

sd 里 cuda内存分配器 pytorch 的区别 gpu内存分配_python_11

如果我现在有两块显卡,我现在都限制其只使用1024MiB,那么我们这样写

sd 里 cuda内存分配器 pytorch 的区别 gpu内存分配_tensorflow_12