pytorch使用GPU炼丹笔记

原理：通过依靠GPU的并行计算能力，能够大大缩短模型训练时间。
在使用GPU跑代码的时候，只需要将模型参数和数据放到GPU上转换成CUDA张量即可。

所以，代码需要修改的地方只有两处：

• 1.模型实例化处。
• 2.数据迭代处(如果数据有更新或增加的话也需要在相应的地方改，只要确定在输入模型前的所有数据都转换为CUDA张量即可)。

使用单GPU

设置设备

device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu") 

cuda:0表示使用0号显卡进行训练，所以如果需要指定其他GPU，则可以直接修改该数字即可。如cuda:1，表示使用1号显卡进行训练。

将数据转换成CUDA张量

也就是将数据放到GPU上

for i,data in enumerate(trainloader,0):
	#从迭代器中获取数据输入
	inputs,labels = data
	inputs, labels = inputs.to(device), labels.to(device) #数据转换为CUDA张量
	outputs=net(inputs) #数据输入模型

to(device)的作用是将数据转换为CUDA张量。

以一个labels为例：

• 加.to(device)之前，labels为tensor([0, 6, 2, 0])
• 加.to(device)之后，labels为tensor([0, 6, 2, 0], device=‘cuda:0’)

将模型参数转换成CUDA张量

同理，也就是将模型放到GPU上

net=net.to(device)  #model为实例化的模型

使用指定GPU

PyTorch默认使用从0开始的GPU，如果GPU0正在运行程序，需要指定其他GPU。

除了上面直接修改cuda:0的方法之外还有以下两种方法来指定需要使用的GPU。

1.使用CUDA_VISIBLE_DEVICES。

1.1 直接在终端或shell脚本中设定：

CUDA_VISIBLE_DEVICES=1 python my_script.py

1.2 python代码中设定：

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2"

2. 使用函数 set_device

import torch
torch.cuda.set_device(id)

该函数见 pytorch-master\torch\cuda_init_.py。
不过官方建议使用CUDA_VISIBLE_DEVICES，不建议使用 set_device 函数。

使用多GPU

DP方法

如果有多个GPU，使用nn.DataParallel来包装我们的模型。 然后通过model.to(device)把模型放到GPU上。
代码如下：

device=torch.device("cuda" if torch.cuda.is_available() else "cpu") 
net=Net().to(device) #需要先加载到GPU上，将模型参数转换成CUDA向量
if torch.cuda.device_count() > 1:
	net=nn.DataParallel(net)

DDP方法

DDP方法比DP方法要好，解决了DP数据分配不平衡和训练速度慢的缺点。
需要添加的地方有：

需要先初始化

需初始化local_rank参数，这里通过argparse模块设置：

parser.add_argument('--local_rank', type=int, default=0,help='node rank for distributed training') 

local_rank参数代表要训练的机子，本来用于标记主机和从机的，如果是多机的话，不同的机器使用不同的local_rank标识机器，由于这里是单机多卡，只用0表示主机就可以了。
另外还需要初始化进程组，代码如下：

torch.distributed.init_process_group(backend="nccl", init_method='file://file:///DATA/wanghongzhi/first_dnn/temp/test',rank=0, world_size=1) 

其中rank是主机的编号，world_size表示分布式主机的个数。

数据集的处理

这步主要是保证一个batch里的数据被均摊到进程上，每个进程能获取到相应的数据。每个GPU产生一个进程。

train_sampler = torch.utils.data.distributed.DistributedSampler(trainset)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,shuffle=False,sampler=train_sampler)

模型初始化

先执行net=Net().to(device)将模型放在GPU上后，再执行下面语句初始化模型：

net = torch.nn.parallel.DistributedDataParallel(net, device_ids=[args.local_rank],output_device=args.local_rank)

单节点多GPU分布式训练

运行代码的shell脚本为：

python -m torch.distributed.launch --nproc_per_node=2  main.py

其中nproc_per_node为使用的显卡数量。

关于DDP的更多细节、参数的选择及其作用可以参考：
博客1
博客2
博客3
博客4
github资源

实验结果

显卡使用TITIAN Xp，从下面数据结果看到DDP效果最好。

• 1.使用CPU耗时888.595s
• 2.使用一块GPU耗时80.06s
• 3.使用DP,两块GPU耗时60.7744s
• 4.使用DDP,两块GPU耗时48.7986s

参考：
单GPU
多GPU