如何使用GPU训练/测试模型
- 使用单GPU
- 设置设备
- 将数据转换成CUDA张量
- 将模型参数转换成CUDA张量
- 使用指定GPU
- 1.使用CUDA_VISIBLE_DEVICES。
- 1.1 直接在终端或shell脚本中设定：
- 1.2 python代码中设定：
- 2. 使用函数 set_device
- 使用多GPU
- DP方法
- DDP方法
- 需要先初始化
- 数据集的处理
- 模型初始化
- 单节点多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 = datainputs, 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
函数)
