dataset

class RandomDataset(Dataset):

def __init__(self, size, length):
    self.len = length
    self.data = torch.randn(length, size)

def __getitem__(self, index):
    return self.data[index]

def __len__(self):
    return self.len

model define

class Model(nn.Module): # Our model

def __init__(self, input_size, output_size):
    super(Model, self).__init__()
    self.fc = nn.Linear(input_size, output_size)

def forward(self, input):
    output = self.fc(input)
    print("\tIn Model: input size", input.size(),
          "output size", output.size())

    return output

if name=="main": # Parameters input_size = 5 output_size = 2

batch_size = 30
data_size = 100

dataset = RandomDataset(input_size, data_size)
# dataloader define
rand_loader = DataLoader(dataset=dataset,
                        batch_size=batch_size, shuffle=True)

# model init
model = Model(input_size, output_size)

# cuda devices
os.environ["CUDA_VISIBLE_DEVICES"]="0,1"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

if torch.cuda.device_count() > 1:
    print("Let's use", torch.cuda.device_count(), "GPUs!")
    # dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs
    model = nn.DataParallel(model)

model.to(device)

for data in rand_loader:
    input = data.to(device)
    output = model(input)
    # loss

    # backward

    #update
    
    time.sleep(1)#模拟一个比较长的batch时间
    print("Outside: input size", input.size(),
        "output_size", output.size())

torch.save(model.module.state_dict(), "model.pth")


</details>

- 如果使用一块GPU，则测试结果为如下，可以看出模型内部与外部输入输出是一致的。


    In Model: input size torch.Size([30, 5]) output size torch.Size([30, 2])

Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2]) In Model: input size torch.Size([30, 5]) output size torch.Size([30, 2]) Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2]) In Model: input size torch.Size([30, 5]) output size torch.Size([30, 2]) Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2]) In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2]) Outside: input size torch.Size([10, 5]) output_size torch.Size([10, 2])

- 如果使用两块GPU，则测试结果如下，可以看出自动进行batch的拆分。


    In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])
    In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])

Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2]) In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2]) In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2]) Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2]) In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2]) In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2]) Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2]) In Model: input size torch.Size([5, 5]) output size torch.Size([5, 2]) In Model: input size torch.Size([5, 5]) output size torch.Size([5, 2]) Outside: input size torch.Size([10, 5]) output_size torch.Size([10, 2])

### 注意
- 多个GPU有主次之分，主GPU负责梯度的整合和参数的更新，将更新后的参数传递给其他GPU，完成一次迭代，所以这个过程中GPU之间既有参数传递，又有参数的梯度传递。
- 关于Batch Norm的问题，[参考](https://discuss.pytorch.org/t/how-does-dataparallel-handels-batch-norm/14040)。因为一个大的batch被拆分成了多个minibatch，所以normalization只会在minibatch上做，最后测试时用的Noramlization layer的参数是主GPU上的。如果想实现多个GPU上同步Normalization需要用[sync norm](https://github.com/vacancy/Synchronized-BatchNorm-PyTorch)实现。这个normalization的问题在分布式训练中依然存在。

## 方案二
方案二是用多进程来实现的，其实分布式就是多进程的意思，分布在多个机器上的进程，利用网络通信协调彼此。关于分布式的处理下一篇文章再详细介绍。这里主要介绍单机上方案二与方案一的不同。首先每个进程都有独立的训练过程，一次迭代后share梯度，整合梯度，独立更新参数。迭代过程中不会进行参数的传递（初始化时会同步所有进程上的参数）。其次进程之间的通信采用了NCCL，当然NCCL已经是pytorch内部支持了，所以一般情况下不用理这个。分布式的细节参考下一篇文章，这里只给出最简单的实现。

### 步骤
- 需要先初始化进程组，这里采用默认的方式初始化，对于单节点来说这也是最方便的一种初始化方式，初始化的目的是让所有的进程彼此建立联系，即知道彼此的位置，状态等信息。
- dataset prepare，增加`torch.utils.data.distributed.DistributedSampler`. 具体使用参见测试部分的代码。
- model prepare, 增加`torch.nn.parallel.DistributedDataParallel`. 具体使用参见测试部分的代码。
- 训练过程与方案一一致，想象同时有多个进程在同时运行training的代码即可。

### 测试
代码与方案一类似，需要初始化进程组，表示本程序是分布式训练的。多进程的创建通过指定`python -m torch.distributed.launch --nproc_per_node=2 --nnodes=1`来实现的，nnodes为1，因为这里我们是一个计算节点，`nproc_per_node=2`表示需要创建两个进程来训练，然后每个进程都获得分配给它rank号，rank唯一标识一个进程，rank 0为master，其他是slave。当然一般是需要两个GPU的，测试程序中是根据rank来指定进程使用GPU，即rank 0使用GPU0，rank 1进程使用GPU1。需要根据数据集创建一个分布式的sampler，初始化dataloader的时候要指定这个sampler，模型分布式封装详见代码。
代码：
<details> 
<summary>展开</summary>  
```python
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import os
import torch.distributed as dist
import torch.utils.data.distributed
import sys
import time


# dataset
class RandomDataset(Dataset):

    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len

# model define
class Model(nn.Module):
    # Our model

    def __init__(self, input_size, output_size):
        super(Model, self).__init__()
        self.fc = nn.Linear(input_size, output_size)

    def forward(self, input):
        output = self.fc(input)
        # print("\tIn Model: input size", input.size(),
        #       "output size", output.size())

        return output

if __name__=="__main__":
    # Parameters
    input_size = 5
    output_size = 2

    batch_size = 30
    data_size = 100

    # check the nccl backend
    if not dist.is_nccl_available():
        print("Error: nccl backend not available.")
        sys.exit(1)

    # init group
    dist.init_process_group(backend="nccl", init_method="env://")

    # get the process rank and the world size
    rank = dist.get_rank()
    world_size = dist.get_world_size()

    # prepare the dataset
    dataset = RandomDataset(input_size, data_size)
    train_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
    

    rand_loader = DataLoader(dataset, batch_size=batch_size//world_size, 
                              shuffle=(train_sampler is None), 
                              sampler=train_sampler)

    # dataloader define
    # rand_loader = DataLoader(dataset=dataset,
    #                         batch_size=batch_size, shuffle=True)

    # model init
    model = Model(input_size, output_size)

    # cuda devices
    # os.environ["CUDA_VISIBLE_DEVICES"]="0"
    # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    # if torch.cuda.device_count() > 1:
    #     print("Let's use", torch.cuda.device_count(), "GPUs!")
    #     # dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs
    #     model = nn.DataParallel(model)
    # model.to(device)

    # distribute model define
    device = torch.device('cuda', rank)
    model = model.to(device)
    model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[rank], output_device=rank)
    print("From rank %d: start training, time:%s"%(rank, time.strftime("%Y-%m-%d %H:%M:%S")))
    for data in rand_loader:
        input = data.to(device)
        output = model(input)
        # loss

        # backward

        #update
        
        time.sleep(1)#模拟一个比较长的batch时间
        print("From rank %d: Outside: input size %s, output size %s"%(rank, str(input.size()), str(output.size())),flush=True)
    torch.save(model.module.state_dict(), "model_%d.pth"%rank)
    print("From rank %d: end training, time: %s"%(rank, time.strftime("%Y-%m-%d %H:%M:%S")))