如果是单个GPU或CPU可通过torch.cuda.is_available()
来设置是否使用GPU
# 单GPU或者CPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
对于多GPU设备而言,如果要定义设备,则需要使用torch.nn.DataParallel
来将模型加载到多个GPU中
# 判断GPU设备大于1,也就是多个GPU,则使用多个GPU设备的id来加载
if torch.cuda.device_count() > 1:
model = nn.DataParallel(model,device_ids=[0,1,2])
也可以使用DistributedDataParallel
分布式计算实现多机多卡的计算。
在pytorch中可以通过.cuda()
以及.to(device)
两种方式将模型和数据复制到GPU的显存中计算。以下说明是在单GPU或者CPU中的情况。
.to(device)
方法可以指定CPU或者GPU。也就是可以定义设备,然后将数据或者模型加载到设备中,代码如下:
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 或device = torch.device("cuda:0")
# 或device = torch.device("cuda:1")
print(device)
# 需要提前初始化一下model
model.to(device)
# 计算的时候将数据加载到device
for batch_idx, (img, label) in enumerate(train_loader):
img=img.to(device)
label=label.to(device)
.cuda()
只能通过os.environ['CUDA_VISIBLE_DEVICE']
来指定GPU
#指定某个GPU
os.environ['CUDA_VISIBLE_DEVICE']='1'
# 判断GPU是否能用,如果能用,则将数据和模型加载到GPU上
if torch.cuda.is_available():
data = data.cuda()
model = model.cuda()
为了加速训练过程,往往会使用多个GPU设备来进行并行训练,常见的多GPU应用场景如下:
在Pytorch中,实现多GPU的方法有torch.nn.DataParallel 和torch.nn.parallel.DistributedDataParallel :
torch.nn.DataParallel
是单进程多线程的方法,仅能工作在单机多卡的情况下。torch.nn.parallel.DistributedDataParallel
方法是多进程多线程的,适用于单机多卡和多机多卡的情况(在单机多卡的情况下DistributedDataParallell
也比DataParallel
的速度更快)在pytorch中,单机多卡和多级多卡的训练教程可参考:PyTorch Distributed Overview
参考:torch.nn.DataParallel Optional: Data Parallelism教程
torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)
该方法的优势在于实现简单,不涉及多进程操作,主要在于使用torch.nn.DataParallel
将模型wrap(包装)一下:
>>> net = torch.nn.DataParallel(model, device_ids=[0, 1, 2])
>>> output = net(input_var) # input_var可以在任何设备上,包括CPU上
该代码实现过程是:首先在前向传播过程中,将输入的数据划分为多个子部分送到不同的device中进行计算,而模型是在每个GPU上都复制一份。也就是说将数据分为device个的子数据输入到各个
device_ids
中,然后将模型复制加载到每个device中对该GPU上的数据进行训练,训练结束后,将每个GPU上计算出来的数据收集到output_device
上合并计算。概括来说就是:DataParallel 会自动帮我们将数据切分并 load 到相应 GPU上,并在每个GPU上都加载模型进行正向传播计算梯度,并汇总到output_device。
可以看出,torch.nn.DataParallel
方法不改变模型的输入和输出,因此只需要将模型用函数torch.nn.DataParallel
将模型wrap(包装)一下即可,无需其他操作。非常方便。但是后续的loss计算只会在output_device上进行,没法并行,因此会导致负载不均衡的问题。并且,对于后续的loss计算,由于来自与多个GPU中,因此需要取一下平均(多维的):
if gpu_count>1:
loss = loss.mean()
其实在使用torch.nn.DataParallel
这个方法使用多GPU的时候,一般可以设置os.environ['CUDA_VISIBLE_DEVICES']
来限制使用的GPU个数:
# 使用编号第0个和第2个的GPU
os.environ['CUDA_VISIBLE_DEVICES'] = '0,2'
需要注意的是,该该参数设置需要在模型加载到GPU之前,一般是在程序最开始,导包之后设定。一般加载方式如下:
# 加载模型
model = torch.nn.DataParallel(model)
model.cuda()
# 加载数据
inputs = inputs.cuda()
labels = labels.cuda()
在官方教程中,加载模型和加载数据方式不是使用.cuda()
方式,但是基本原理差不过,查看DataParallel
内部实现发现其实差不多。
由于DataParallel
会存在负载不均衡问题。官方提供的解决方案是使用torch.nn.parallel.DistributedDataParallel来代替。而且DistributedDataParallel运行更快,现存分配更加均衡,功能更加强悍。
使用DistributedDataParallel
(简称DDP)来进行分布式计算需要考虑:
而DDP使用流程如下:
首先可通过os.environ['CUDA_VISIBLE_DEVICES']
来设置指定的GPU来使用。然后初始化:
torch.distributed.init_process_group(backend='nccl', init_method='tcp://localhost:23456', rank=0, world_size=1)
- backend:后端,实际上是多个机器之间交换数据的协议。用于分布式训练时,多个计算设备之间的集合通信,常见的有
Open MPI
(不怎么用)、NCCL
、Gloo
等。后端选择好了之后, 因为多个主机之间需要通过网络进行数据交换,需要设置网络接口, 对于nccl
和gloo
一般会自己寻找网络接口, 但是某些时候, 可能需要自己手动设置:import os # 说明,其中网卡名字需要自己去找,如果是linux可通过ifconfig来查看。 # 使用gloo后端需要设置的 os.environ['GLOO_SOCKET_IFNAME'] = 'eth0' # 使用nccl需要设置的 os.environ['NCCL_SOCKET_IFNAME'] = 'eth0'
init_method
:机器之间交换数据, 需要指定一个主节点, 而这个参数就是指定主节点的。有TCP
和共享文件系统
等方式设置。使用TCP
格式为tcp://ip:端口号
,ip是主节点的ip地址,rank参数为0的那个主机ip地址,然后选择一个空闲的端口号即可。
rank
:rank是标识主机和从机的, 主节点为0, 剩余的为了1-(N-1), N为要使用的机器的数量。
world_size
介绍说是进程个数,其实就是机器的个数,如果有两台机器一起训练的话,world_size就是2
模型的处理其实和上面的单机多卡没有多大区别, 还是下面的代码, 但是注意要提前想把模型加载到gpu, 然后才可以加载到DistributedDataParallel
model = model.cuda()
model = nn.parallel.DistributedDataParallel(model)
对于torch.utils.data.distributed.DistributedSampler
方法的作用是将数据集进行打乱,然后根据GPU的个数自行补充数据,并对补充后的数据分配到不同的GPU设备上。如图:
对于torch.utils.data.BatchSampler
方法的作用是将torch.utils.data.distributed.DistributedSampler
分配给GPU的数据按照batch_size分配成一个个的batch。
# 实例化训练数据集
train_data_set = MyDataSet(images_path=train_images_path,
images_class=train_images_label,
transform=data_transform["train"])
# 实例化验证数据集
val_data_set = MyDataSet(images_path=val_images_path,
images_class=val_images_label,
transform=data_transform["val"])
# 给每个rank对应的进程分配训练的样本索引
train_sampler = torch.utils.data.distributed.DistributedSampler(train_data_set)
val_sampler = torch.utils.data.distributed.DistributedSampler(val_data_set)
# BatchSampler是对train_sampler做进一步的处理,组成一个一个的batch
train_batch_sampler = torch.utils.data.BatchSampler(train_sampler,
batch_size,
drop_last=True)
val_batch_sampler = torch.utils.data.BatchSampler(val_sampler,
batch_size,
drop_last=True)
在Pytorch中使用多GPU的常用启动方式一种是
torch.distributed.launch
一种是torch.multiprocessing
模块。这两种方式各有各的好处,在我使用过程中,感觉torch.distributed.launch启动方式更方便,而且我看官方提供的多GPU训练FasterRCNN源码就是使用的torch.distributed.launch方法,所以我个人也比较喜欢这个方法。但在官方的教程中主要还是使用的torch.multiprocessing方法,官方说这种方法具有更好的控制和灵活性。在自己使用体验过程中确实和官方说的一样。
这里提醒下要使用torch.distributed.launch启动方式的小伙伴。训练过程中如果你强行终止的程序,在开启下次训练前建议你通过nvidia-smi指令看下你GPU的显存是否全部释放了,如果没有全部释放,需要手动杀下进程。在我使用过程中发现强行终止程序有小概率出现进程假死的情况,占用的GPU的资源并没有及时释放,如果在下次训练前没有及时释放,会影响你的训练,或者直接提示通信端口被占用,无法启动的情况。
# 使用SyncBatchNorm后训练会更耗时
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
当然以上总结是我遇到的问题,没有在多机多卡上调试过,主要介绍多机多卡是因为torch.nn.parallel.DistributedDataParallel
在很多开源代码中使用了,因为该方法可以在单机单卡上很好运行,但是在调试代码的时候可能会出现一些关于torch.nn.parallel.DistributedDataParallel
方法的错误。这里简单介绍一下。后续如果遇到问题,可能还会继续补充。
补充知识点:
为什么要使用os.environ
?
在python中,environ是一个字符串所对应环境的映像对象,怎么理解呢,就是可以通过os.environ
获取有关系统的各种信息(比如环境变量),常用用法如下:
# 设置
os.environ['环境变量名称']='环境变量值' #其中key和value均为string类型
os.putenv('环境变量名称', '环境变量值')
os.environ.setdefault('环境变量名称', '环境变量值')
# 修改
os.environ['环境变量名称']='新环境变量值'
# 获取
os.environ['环境变量名称']
os.getenv('环境变量名称')
os.environ.get('环境变量名称', '默认值') #默认值可给可不给,环境变量不存在返回默认值
# 删除
del os.environ['环境变量名称']
del(os.environ['环境变量名称'])
# 判断
'环境变量值' in os.environ # 存在返回 True,不存在返回 False
在linux和window中存在一些常见的环境变量名称,这部分可以自行去查看API,对于pytorch中,使用os.environ['CUDA_VISIBLE_DEVICES'] = '0,2'
表示指定使用GPU。
参考:
NCCL、OpenMPI、Gloo对比_taoqick的博客-CSDN博客_gloo nccl