torch教程——数据并行处理

pytorch的优点之一是可以使用GPU加速，但是默认情况下，只会使用一个GPU。本文内容是如果使用多个GPU并行运算。

代码

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader

input_size = 5
output_size = 2

batch_size = 30
data_size = 100

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

#数据集：size是单个向量的大小，length是向量的个数
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

#加载器：batch_size = 30，即单次循环中读入的数据是30 * 5
rand_loader = DataLoader(dataset=RandomDataset(size=input_size, length=data_size), batch_size = batch_size, shuffle=True)

#神经网络：只有一个线性层，输入5，输出2
class Model(nn.Module):

    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("\t In Model: input size", input.size(), "output Size", output.size())

        return output


#实例化model
model = Model(input_size, output_size)

#判断是否存在多个GPU
if torch.cuda.device_count() > 1:
    print("Let's use", torch.cuda.device_count(), "GPUS")
    #将模型部署到多个GPU上
    model = nn.DataParallel(model)

#神经网络移植到GPU
model.to(device)

#开始训练
for data in rand_loader:
	#输入移植到GPU
    input = data.to(device)
    output = model(input)
    print("Outsida:input size", input.size(), "output_size", output.size())

这里并行加速的核心代码是

model = nn.DataParallel(model)

数据并行自动拆分了数据并且将任务发送到多个GPU上，当每个模型都完成自己的任务后，DataParallel收集并合并这些结果并返回。

输出

此处，我有两个可用的GPU。
每次for循环中读取30个长度是5的向量，由于有2个GPU，因此每个GPU分配15个，在模型内部，就是15*5的张量。

说明

1. 使用多个GPU并行处理速度不一定更快。笔者测试的数据集较小，神经网络也很简单，使用多个GPU训练用时可能不降反升。