PyTorch的GPU训练方式学习

GPU并行训练

模型转cuda

gpus=[0]#指定训练的GPU
cuda_gpu=torch.cuda.is_available()   #判断GPU是否存在可用

net=Net()#模型初始化
if(cuda_gpu):
	net=torch.nn.DataParallel(net,device_ids=gps).cuda()#ids指定GPU

数据转cuda

(minibatchX, minibatchY) = minibatch
minibatchX = minibatchX.astype(np.float32).T#转浮点型数字
minibatchY = minibatchY.astype(np.float32).T
if(cuda_gpu):
    b_x = Variable(torch.from_numpy(minibatchX).cuda())    #将数据转为cuda类型
    b_y = Variable(torch.from_numpy(minibatchY).cuda())
else:
    b_x = Variable(torch.from_numpy(minibatchX))
    b_y = Variable(torch.from_numpy(minibatchY))

输出数据去cuda，转为numpy

correct_prediction = sum(torch.max(output, 1)[1].data.squeeze() == torch.max(b_y, 1)[1].data.squeeze())
if(cuda_gpu):
    correct_prediction = correct_prediction.cpu().numpy()   #.cpu将cuda转为tensor类型，.numpy将tensor转为numpy类型
else:
    correct_prediction = correct_prediction.numpy()

单机多卡 分布式训练

补充文章：DDP入门教程
DDP实现原理与源代码解析
DDP实战技巧

torch.distributed 包支持

分布训练包括 GPU 和 CPU 的分布式训练支持。Pytorch 分布式目前只支持 Linux。
torch.distributed 分布式训练的优势：

1.每个进程对应一个独立的训练过程，且只对梯度等少量数据进行信息交换。

• 在每次迭代中，每个进程具有自己的 optimizer ，并独立完成所有的优化步骤，进程内与一般的训练无异。
• 在各进程梯度计算完成之后，各进程需要将梯度进行汇总平均，然后再由 rank=0 的进程，将其 broadcast(广播) 到所有进程。之后，各进程用该梯度来更新参数
• 由于各进程中的模型，初始参数一致 (初始时刻进行一次 broadcast)，而每次用于更新参数的梯度也一致，因此，各进程的模型参数始终保持一致
  而在 DataParallel 中，全程维护一个 optimizer，对各 GPU 上梯度进行求和，而在主 GPU 进行参数更新，之后再将模型参数 broadcast 到其他 GPU。所以分布式训练的参数传输的数据量更少，因此速度更快效率更高。
  基本概念
• group
  即进程组。默认情况下，只有一个组，一个 job 即为一个组，也即一个 world。当需要进行更加精细的通信时，可以通过 new_group 接口，使用 word 的子集，创建新组，用于集体通信等。
• world size
  全局进程个数
• rank
  进程序号，用于进程间通讯，表征进程优先级。rank=0的主机为mask
• local_rank
  进程内GPU编号，非显式参数，由 torch.distributed.launch 内部指定。比方说， rank = 3，local_rank = 0 表示第 3 个进程内的第 1 块 GPU。
  基本使用流程

1. 在使用 distributed 包的任何其他函数之前，需要使用 init_process_group 初始化进程组，同时初始化 distributed 包

distributed.init_process_group(backend='nccl',rank,world_size,init_method='env://')
#Env初始化方式

2. 如果需要进行小组内集体通信，用 new_group 创建子分组
3. 创建分布式并行模型 DDP(model, device_ids=device_ids)
4. 为数据集创建 Sampler

train_sample=DistributeedSampler(traindataset)

5. 使用启动工具 torch.distributed.launch --n在每个主机上执行一次脚本，开始训练
6. 使用 destory_process_group() 销毁进程组

编程模板
一,Env初始化方法

import torch.distributed as dist
import torch.utils.data.distributed

# ......
import argparse
parser = argparse.ArgumentParser()
# 注意这个参数，必须要以这种形式指定，即使代码中不使用。因为 launch 工具默认传递该参数
parser.add_argument("--local_rank", type=int)
args = parser.parse_args()

# ......
dist.init_process_group(backend='nccl', init_method='env://')

# ......
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=download, transform=transform)
train_sampler = torch.utils.data.distributed.DistributedSampler(trainset)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, sampler=train_sampler)

# ......
# 根据 local_rank，配置当前进程使用的 GPU
net = Net()
device = torch.device('cuda', args.local_rank)
net = net.to(device)
net = torch.nn.parallel.DistributedDataParallel(net, device_ids=[args.local_rank], output_device=args.local_rank)

执行方式

python -m torch.distributed.launch --nproc_per_node=2 --nnodes=3 --node_rank=0 --master_addr="192.168.1.201" --master_port=23456 env_init.py

python -m torch.distributed.launch --nproc_per_node=2 --nnodes=3 --node_rank=1 --master_addr="192.168.1.201" --master_port=23456 env_init.py

python -m torch.distributed.launch --nproc_per_node=2 --nnodes=3 --node_rank=2 --master_addr="192.168.1.201" --master_port=23456 env_init.py

其中，节点数为3，每个节点的进程数为2，node_rank是每个节点的优先级，–master_addr是指定主节点

说明

1. 在 Env 方式中，在 init_process_group 中，无需指定任何参数
2. 必须在rank==0的进程内保存参数
3. 该方式下，使用 torch.distributed.launch 在每台主机上，为其创建多进程，其中:
4. nproc_per_node 参数指定为当前主机创建的进程数。一般设定为当前主机的 GPU 数量
5. nnodes 参数指定当前 job 包含多少个节点
6. node_rank 指定当前节点的优先级
7. master_addr 和 master_port 分别指定 master 节点的 ip:port
8. 若没有为每个进程合理分配 GPU，则默认使用当前主机上所有的 GPU。即使一台主机上有多个进程，也会共用 GPU
9. 使用 torch.distributed.launch 工具时，将会为当前主机创建 nproc_per_node 个进程，每个进程独立执行训练脚本。同时，它还会为每个进程分配一个 local_rank 参数，表示当前进程在当前主机上的编号。例如：rank=2, local_rank=0 表示第 3 个节点上的第 1 个进程。
10. 需要合理利用 local_rank 参数，来合理分配本地的 GPU 资源
11. 每条命令表示一个进程。若已开启的进程未达到 word_size 的数量，则所有进程会一直等待
    二,TCP初始化方式
    代码：

import torch.distributed as dist
import torch.utils.data.distributed

# ......
parser = argparse.ArgumentParser(description='PyTorch distributed training on cifar-10')
parser.add_argument('--rank', default=0,
                    help='rank of current process')
parser.add_argument('--word_size', default=2,
                    help="word size")
parser.add_argument('--init_method', default='tcp://127.0.0.1:23456',
                    help="init-method")
args = parser.parse_args()

# ......
dist.init_process_group(backend='nccl', init_method=args.init_method, rank=args.rank, world_size=args.word_size)

# ......
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=download, transform=transform)
train_sampler = torch.utils.data.distributed.DistributedSampler(trainset)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, sampler=train_sampler)

# ......
net = Net()
net = net.cuda()
net = torch.nn.parallel.DistributedDataParallel(net)

执行方式：

# Node 1 : ip 192.168.1.201  port : 12345
python tcp_init.py --init_method tcp://192.168.1.201:12345 --rank 0 --word_size 3

# Node 2 : 
python tcp_init.py --init_method tcp://192.168