pytorch指定GPU训练

主要讲单机多卡(单主机多GPUs训练)
使用多卡训练的方式有很多,当然前提是我们的设备中存在多个GPU:使用命令nvidia-smi查看当前Ubuntu平台的GPU数量,其中每个GPU被编上了序号:[0,1,2,3]

###################################################################

1、torch.cuda.is_available() #cuda是否可用

2、torch.cuda.device_count()#GPU 的数量

3、torch.cuda.current_device() #当前设备的索引,从0开始

4、torch.cuda.get_device_name(0)#返回GPU名字
####################################################################

0. 显卡编号(什么是主卡)
在默认情况下,标号为0的显卡为主卡,如主机中有4块显卡,那么每张显卡的默认标号为[0,1,2,3]。

如何将其他显卡设置为主卡呢?

通过os.environ["CUDA_VISIBLE_DEVICES"]指定所要使用的显卡,如:

os.environ["CUDA_VISIBLE_DEVICES"] = "3,2,0,1"
model = torch.nn.DataParallel(model, device_ids=[0,2,3]).cuda()


此时,3号显卡就变成了主卡,在使用torch.nn.DataParallel指定运算显卡时,显卡的对应关系如下:

实际显卡编号----->运算显卡编号
    3     ----->     0(主卡)
    2     ----->     1
    0     ----->     2
    1     ----->     3


1. 使用所有存在的显卡
在存在多卡的条件下,最简单的方法是直接使用torch.nn.DataParallel将你的模型wrap一下即可:

net = torch.nn.DataParallel(model)


这时,默认所有存在的显卡都会被使用

2. 指定编号使用显卡


如果有很多显卡(例如我们有4张显卡),但只想使用0、1、2号显卡,那么可以:

net = torch.nn.DataParallel(model, device_ids=[0, 1, 2])


或者:

os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(map(str, [0,1,2])) # 一般在程序开头设置
# 等价于os.environ["CUDA_VISIBLE_DEVICES"] = '0,1,2'
net = torch.nn.DataParallel(model)


CUDA_VISIBLE_DEVICES 表示当前可以被python环境程序检测到的显卡

如果只需要指定一张卡,可以使用torch.cuda.set_device(1)指定gpu使用编号
(不建议用这个方法)

torch.cuda.set_device(1)
print(torch.cuda.device_count()) #可用GPU数量


(我的机器是4卡,所以print结果是:4,说明用torch.cuda.set_device(1)指定,不会改变可见的显卡)
后面还可以用torch.nn.DataParallel(model, device_ids=[1, 2])进行指定,
但是必须包含set_device(1)指定的device:1的设备,缺点是仍然会存在占用一些device:0的gpu内存;

3. os.environ[“CUDA_VISIBLE_DEVICES”]详解


3.1. [“CUDA_VISIBLE_DEVICES”]使用
CUDA_VISIBLE_DEVICES 表示当前可以被python环境程序检测到的显卡

os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(map(str, [0,1,2])) # 一般在程序开头设置
# 等价于os.environ["CUDA_VISIBLE_DEVICES"] = '0,1,2'


os.environ["CUDA_VISIBLE_DEVICES"] = '1,2'进行指定使用设备
这样会修改pytorch感受的设备编号,pytorch感知的编号还是从device:0开始。如上则把1号显卡改为device:0,2号显卡改为device:1,使用时应该这么写:

os.environ["CUDA_VISIBLE_DEVICES"] = '1,2'
torch.nn.DataParallel(model, device_ids=[0,1])



3.2. 关于设置[“CUDA_VISIBLE_DEVICES”]无效的解决
不生效的原因是,这一行代码放置的位置不对。
一定要把os.environ['CUDA_VISIBLE_DEVICES']='1'放在所有访问GPU的代码之前。

torch.cuda.device_count()

4. torch.cuda主要函数
4.1. torch.cuda.is_available()判断GPU是否可用

import torch
print(torch.cuda.is_available())


可用即为:True

4.2. torch.cuda.device_count()查看可用GPU数量

os.environ["CUDA_VISIBLE_DEVICES"] = '0,1,2'
print(torch.cuda.device_count())


结果应该为:3

4.3. torch.cuda.current_device()查看当前使用的GPU序号

os.environ["CUDA_VISIBLE_DEVICES"] = '1,2'
print(torch.cuda.current_device())


结果应该是:0
这样就证明了我们上面说的,os.environ["CUDA_VISIBLE_DEVICES"] = '1,2'进行指定使用设备,会修改pytorch感受的设备编号,pytorch感知的编号还是从device:0开始

4.4 其他一些函数


查看指定GPU的容量、名称

torch.cuda.get_device_capability(device)
torch.cuda.get_device_name(device)



清空程序占用的GPU资源

torch.cuda.empty_cache()


为GPU设置随机种子

torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)


5. 报错AssertionError: Invalid device id


对模型进行gpu指定编号多gpu训练,必须要有所指定编号的gpu,不然会报AssertionError: Invalid device id错误

报错示例1:

os.environ["CUDA_VISIBLE_DEVICES"] = '0,1' # 一般在程序开头设置
net = torch.nn.DataParallel(model,device_ids=[1, 2])

CUDA_VISIBLE_DEVICES 表示当前可以被python环境程序检测到的显卡,可见的只有0,1号显卡。
而使用torch.nn.DataParallel(model, device_ids=[1, 2])指定gpu编号会出现AssertionError: Invalid device id错误
原因:2号显卡没有设置被可见。

报错示例2:

os.environ["CUDA_VISIBLE_DEVICES"] = '1,2' # 一般在程序开头设置
net = torch.nn.DataParallel(model,device_ids=[1, 2])


os.environ["CUDA_VISIBLE_DEVICES"] = '1,2'进行指定使用设备
这样会修改pytorch感受的设备编号,pytorch感知的编号还是从device:0开始
程序应该这么改:

os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(map(str, [1,2])) # 一般在程序开头设置
net = torch.nn.DataParallel(model,device_ids=[0, 1])


6. 报错RuntimeError: module must have its parameters and buffers on device cuda:1 (device_ids[0]) but found
PyTorch框架下,使用torch.nn.DataParallel进行多卡并行计算中可能会遇到该问题,是由于在多卡运算时主卡设置所导致。在多卡并行计算时,所使用的显卡中必须包含主卡,否则就会产生该问题。

如下代码中,主卡没有参与运算,就会导致该错误。

os.environ["CUDA_VISIBLE_DEVICES"] = '0,1,2,3' #这里主卡就是机器0号卡
model = torch.nn.DataParallel(model, device_ids=[1,2,3]).cuda()


————————————————
版权声明:本文为CSDN博主「夏普通」的原创文章,遵循CC 4.0 BY-SA版权协议,转载请附上原文出处链接及本声明。
原文链接:https://blog.csdn.net/qq_34243930/article/details/106695877

你可能感兴趣的:(pytorch,python)