我理解的分布式训练主要是通过两种方式实现：

1）数据并行

2）模型并行

 

接下来，本文从3个方面来简单介绍一下分布式训练：

1）单个GPU

2）单机多卡

3）多机多卡（分布式训练）

 

一、单机单卡

PyTorch集成了一套高效的GPU运算体系，可以加快数据运算，大概比CPU要快几十倍~

先上代码：

# 方法介绍
device = torch.device(param)  # param可以是'cuda'或者'cpu'，该方法返回一个对象，表示数据将要放置的设备
device = torch.device('cuda',0) # 指定GPU设备，0为设备号
torch.cuda.is_available()      # 查看当前GPU设备是否可用
#-------------------------------------- start --------------------------------------------
# 简单使用说明
# 在PyTorch上使用GPU进行训练非常简单（三行代码）
# 首先，看设备是否可用
import torch
x = torch.rand(3,4) # 随机生成一个3行4列的张量x
if torch.cuda.is_available():   #返回True，则GPU可用
    # 再者，指定GPU设备对象
    device = torch.device('cuda') # 或者 device = torch.device('cuda:0') 或者 device = torch.device('cuda',0) 都可，因为是单卡，所以为0

# 最后将张量x移到GPU上去，有两种方法，这两种方法可以认为是等价的，无明显区别
x = x.to(device) # 方法1
x = x.cuda()     # 方法2
#--------------------------------------- end ---------------------------------------------

上述代码可以完成一件事：将数据送上GPU；除此之外，可能需要注意以下几点：

1、模型上GPU：model.cuda()   

2、数据上GPU：data.cuda()   

3、输出下GPU：output = model(data)  output.detach().cpu().numpy()，

     ===> .detach()的作用是将变量output从计算图中分离，使其不具有梯度，不进行反向传播。.cpu()是将GPU数据转CPU，.numpy()是将Tensor转numpy，如果需要继续反向传播，则不需要.detach().cpu().numpy()。一般如果要单独处理，如图像的显示，或者特征的保存，都需要做上述操作。

二、单机多卡

目前PyTorch的单机多卡训练，主要有两种方式：

# 第一种
torch.nn.DataParallel()
# 第二种
torch.nn.parallel.DistributedDataParallel()

其中，第一种只能在单机多卡模式下训练；第二种可以实现单机多卡和多机多卡，真正实现分布式训练。除此之外，性能上，第二种方法优于第一种。下面说怎么用：

# 第一种 torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)
# 方法介绍
torch.cuda.device_count() # 返回当前设备可用的GPU数量
#-------------------------------------- start ----------------------------------------
# 使用方法(假设网络模型为model)
import torch
x = torch.rand(3,4)
model = torch.nn.DataParallel(model,device_ids=list(range(torch.cuda.device_count())) # 主要就是这句
output = model(x)
#--------------------------------------  end  ----------------------------------------

其中，device_ids为指定设备号，为列表，默认为所有设备；output_device为输出的设备号，默认是0号设备；经过主要的那句，就可以实现单机多卡训练。注意：此时的 batch_size 应该是 单机单卡的bach_size * n，n = torch.cuda.device_count()。该方法默认从数据的第一维进行分割，所以第一维应该是batch_size，如图像数据，数据维度应该是(batch_size,c,h,w)。

# 第二种 torch.nn.parallel.DistributedDataParallel(module, device_ids=None,output_device=None) 参数很多，我这里只给出了三个主要的参数
# 方法介绍
local_rank = torch.distributed.get_rank()  # 返回当前进程组的排名,rank是分布式进程组中每个进程的唯一id。
torch.cuda.set_device(i) # 配置i进程的GPU
os.environ['SLURM_NTASKS']          #可用作world size
os.environ['SLURM_NODEID']          #node id
os.environ['SLURM_PROCID']          #可用作全局rank
os.environ['SLURM_LOCALID']         #local_rank
os.environ['SLURM_STEP_NODELIST']   #从中取得一个ip作为通讯ip
#------------------------------------ start ------------------------------------------
# 首先，初始化
import os
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.distributed import DistributedSampler
from torch.nn.parallel import DistributedDataParallel

torch.multiprocessing.set_start_method('spawn')
rank = int(os.environ['SLURM_PROCID'])
local_rank = int(os.environ['SLURM_LOCALID'])
world_size = int(os.environ['SLURM_NTASKS'])
ip = get_ip(os.environ['SLURM_STEP_NODELIST']) # 确保所有任务的获取的ip一致 （自己实现）
host_addr = 'tcp://' + ip + ':' + str(port)  # 主机的ip地址和端口号 （自己配置到ip和端口）
torch.distributed.init_process_group(backend="nccl", init_method=host_addr, rank=rank, world_size=world_size) # 注：这里的backend表示通信后端，非常多，如NCCL、MPI、Gloo，最快的应该还是nccl
num_gpus = torch.cuda.device_count()
torch.cuda.set_device(local_rank)
assert torch.distributed.is_initialized()
device = torch.device('cuda',local_rank) # 配置GPU设备

# 获取数据和采样器
dataset = your_dataset()  
datasampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
dataloader = DataLoader(dataset, batch_size=batch_size_per_gpu, sampler=source_sampler)

# 配置模型
model = your_model()     #也是按自己的模型写
model.to(device)  # 先把模型放到当前进程的GPU中去
model = DistributedDataPrallel(model, device_ids=[local_rank], output_device=local_rank)
if num_gpus >1:
    model = torch.nn.parallel.DistributedDataParallel(model,device_ids=[local_rank],output_device=local_rank) # 再将模型分发到其他各个GPU

# 此后训练就应该与之前的无异了
# 运行时输入以下指令，这里这么设置，是为了让os.environ获取到相应的信息
srun -n 8 --gres=gpu:4 --ntasks-per-node=4 python train.py # 单机多卡，一台机子，共4张GPU
srun -n 12 --gres=gpu:4 --ntasks-per-node=4 python train.py # 多级多卡，3台设备，每台4张GPU

# ---------------------------------------------------------------------------------------------------------------------------------------

还没写完，之后再补充吧