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

1. 導入和參數

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例子：

虛擬數據集:制作一個虛擬(隨機)數據集, 只需要實現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)

構建一個簡單模型

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

★核心

創建模型對象和數據并行,此部分是本講核心.
首先，需要創建一個模型實例對象，并且檢測是否由多個GPU, 如果多個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)
#)

運行模型
現在可以看到只有1個GPU時候，對30個輸入輸出的處理。

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個GPU，你看到的輸出如下：

image.png

如果有3個GPU, 你會看到輸出如下:

image.png

image.png

總結

DataParallel會自動劃分上數據，并將作業發送到多個GPU上的Duo個模型，并且在每個模型完成作業后，收集合并結果并返回。