這一講主要介紹如何使用DataParallel使用多GPU，

1. 導(dǎo)入和參數(shù)

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

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

以下為多GPU例子：

虛擬數(shù)據(jù)集:制作一個(gè)虛擬(隨機(jī))數(shù)據(jù)集, 只需要實(shí)現(xiàn)getitem

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
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)

構(gòu)建一個(gè)簡(jiǎn)單模型

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

★核心

創(chuàng)建模型對(duì)象和數(shù)據(jù)并行,此部分是本講核心.
首先，需要?jiǎng)?chuàng)建一個(gè)模型實(shí)例對(duì)象，并且檢測(cè)是否由多個(gè)GPU, 如果多個(gè)GPU就是用nn.DataParallel來包裝模型，然后通過model.to(device)把模型放到GPU上。

model = Model(input_size, output_size)
device = torch.device('cuda:0' 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)
#Model(
#   (fc): Linear(in_features=5, out_features=2, bias=True)
#)

運(yùn)行模型
現(xiàn)在可以看到只有1個(gè)GPU時(shí)候，對(duì)30個(gè)輸入輸出的處理。

for data in rand_loader:
    input = data.to(device)
    output = model(input)
    print("Outside: input size", input.size(),
          "output_size", output.size())

image.png

如果有2個(gè)GPU，你看到的輸出如下：

image.png

如果有3個(gè)GPU, 你會(huì)看到輸出如下:

image.png

image.png

總結(jié)

DataParallel會(huì)自動(dòng)劃分上數(shù)據(jù)，并將作業(yè)發(fā)送到多個(gè)GPU上的Duo個(gè)模型，并且在每個(gè)模型完成作業(yè)后，收集合并結(jié)果并返回。