Ch01-04.数据并行(Pytorch入门)

这一讲主要介绍如何使用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())