Pytorch并行分布式训练二 （ DP）

Optional: Data Parallelismhttps://pytorch.org/tutorials/beginner/blitz/data_parallel_tutorial.html#optional-data-parallelism

DP是最简单的pytorch多gpu训练方法，虽然现在已不是pytorch官方的推荐，但也值得做适当的了解。

1.单GPU训练

可以概况为：

1. Model to GPU
2. Input to GPU
3. Input to Model

值得注意的是：

nn.Module实现的to方法返回值为'self',无需重写变量。

nn.Tensor实现的to方法返回值为一个新的Tensor量，注意变量需要重写，否则操作无效。

2.Data Parallel的多GPU训练

 可以概况为：

1. Model to DataParallel （只需要多加这一个步骤）
2. Model to GPU
3. Input to GPU
4. Input to Model

值得注意的是：

torch.nn.DataParallel(module, device_ids=None, output_device=None …)

device_ids参数可以指定需要用到的特定GPU列表，参数的类型为list[int] 或 list[torch.device], 默认为None时，表示使用机器中所有可见的GPU。

 

3.Data Parallel的多GPU训练实例

 

## Imports and parameters

import torch

import torch.nn as nn

from torch.utils.data import Dataset, DataLoader

# Parameters and DataLoaders

input_size = 5

output_size = 2

batch_size = 30

data_size = 100

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

 

## 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

rand_loader = DataLoader(dataset=RandomDataset(input_size, data_size),

                         batch_size=batch_size, shuffle=True)

 

## Model

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

 

##[important] Create Model and DataParallel 

model = Model(input_size, output_size)

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)

 

## Run the Model

for data in rand_loader:

    input = data.to(device)

    output = model(input)

    print("Outside: input size", input.size(),

          "output_size", output.size())

输出结果及分析

双GPU

分析及说明：

1. 从连续输出的两次‘In Model xxx’可以看出，在前向传播时模型被复制成了两个。
2. 从input_size=15和outpu_size=30的数量情况可以看出，在前向传播时，输入batch_size=30被平均分到了两个模型分别计算。在模型输出时两个结果进行了整合，一般合并在device_ids[0]号卡上，这个可以通过参数output_device来指定。
3. 此方法实际batch_size即是你在DataLoader中设置的batch_size，模型计算时的batch_size实际为a或b,其中

（a x n_gpus + b）= batch_size  且 a > b