воскресенье, 30 декабря 2018 г.

PyTorch: использование GPU

В этом посте мы рассмотрим, как в PyTorch использовать несколько GPU, с помощью DataParallel.

С PyTorch очень легко использовать GPU. Вы можете поместить модель на GPU:

device = torch.device("cuda:0")
model.to(device)

Затем вы можете скопировать все свои тензоры в GPU:

mytensor = my_tensor.to(device)

Обратите внимание, что простой вызов my_tensor.to(device) возвращает новую копию my_tensor на GPU вместо переписывания my_tensor. Вы должны назначить его на новый тензор и использовать этот тензор на GPU.

Естественно выполнять прямое и обратное распространение на нескольких GPU. Однако по умолчанию Pytorch будет использовать только один GPU. Вы можете легко запустить свои операции на нескольких GPU, заставляя вашу модель работать параллельно, используя DataParallel:

model = nn.DataParallel(model)

Рассмотрим это подробней.

Импорт и параметры

Импортируйте модули PyTorch и определите параметры.

# -*- coding: utf-8 -*-
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader

# Параметры и 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")


Фиктивный набор данных

Создаем фиктивный (случайный) набор данных. Вам только необходимо реализовать getitem

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)


Выполняем модель

Теперь мы можем увидеть размеры вводных и выводных тензоров.

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

Вывод:

In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])
        In Model: input size torch.Size([15, 5]) 
                  output size torch.Size([15, 2])
Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
        In Model: input size torch.Size([15, 5]) 
                  output size torch.Size([15, 2])
        In Model: input size torch.Size([15, 5]) 
                  output size torch.Size([15, 2])
Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
        In Model: input size torch.Size([15, 5]) 
                  output size torch.Size([15, 2])
        In Model: input size torch.Size([15, 5]) 
                  output size torch.Size([15, 2])
Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
        In Model: input size torch.Size([5, 5]) 
                  output size torch.Size([5, 2])
        In Model: input size torch.Size([5, 5]) 
                  output size torch.Size([5, 2])
Outside: input size torch.Size([10, 5]) output_size torch.Size([10, 2])

Результаты

Если у вас нет GPU или один GPU, то в этом случае, когда мы пакетизируем 30 вводов и 30 выводов, модель получает 30 и выводит 30 как ожидается. Но если у нас есть несколько GPU, тогда мы можем получить результаты как следующие.

2 GPU

Если у вас 2 GPU, тогда вы увидите:

# на 2 GPU
Let's use 2 GPUs!
    In Model: input size torch.Size([15, 5]) 
              output size torch.Size([15, 2])
    In Model: input size torch.Size([15, 5]) 
              output size torch.Size([15, 2])
Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
    In Model: input size torch.Size([15, 5]) 
              output size torch.Size([15, 2])
    In Model: input size torch.Size([15, 5]) 
              output size torch.Size([15, 2])
Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
    In Model: input size torch.Size([15, 5]) 
              output size torch.Size([15, 2])
    In Model: input size torch.Size([15, 5]) 
              output size torch.Size([15, 2])
Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
    In Model: input size torch.Size([5, 5]) 
              output size torch.Size([5, 2])
    In Model: input size torch.Size([5, 5]) 
              output size torch.Size([5, 2])
Outside: input size torch.Size([10, 5]) output_size torch.Size([10, 2])

3 GPU

Если у вас 3 GPU, тогда вы увидите:

Let's use 3 GPUs!
    In Model: input size torch.Size([10, 5]) 
              output size torch.Size([10, 2])
    In Model: input size torch.Size([10, 5]) 
              output size torch.Size([10, 2])
    In Model: input size torch.Size([10, 5]) 
              output size torch.Size([10, 2])
Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
    In Model: input size torch.Size([10, 5]) 
              output size torch.Size([10, 2])
    In Model: input size torch.Size([10, 5]) 
              output size torch.Size([10, 2])
    In Model: input size torch.Size([10, 5]) 
              output size torch.Size([10, 2])
Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
    In Model: input size torch.Size([10, 5]) 
              output size torch.Size([10, 2])
    In Model: input size torch.Size([10, 5]) 
              output size torch.Size([10, 2])
    In Model: input size torch.Size([10, 5]) 
              output size torch.Size([10, 2])
Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
    In Model: input size torch.Size([4, 5]) 
              output size torch.Size([4, 2])
    In Model: input size torch.Size([4, 5]) 
              output size torch.Size([4, 2])
    In Model: input size torch.Size([2, 5]) 
              output size torch.Size([2, 2])
Outside: input size torch.Size([10, 5]) output_size torch.Size([10, 2])

8 GPU

Если у вас 8 GPU, тогда вы увидите:

Let's use 8 GPUs!
    In Model: input size torch.Size([4, 5]) 
              output size torch.Size([4, 2])
    In Model: input size torch.Size([4, 5]) 
              output size torch.Size([4, 2])
    In Model: input size torch.Size([2, 5]) 
              output size torch.Size([2, 2])
    In Model: input size torch.Size([4, 5]) 
              output size torch.Size([4, 2])
    In Model: input size torch.Size([4, 5]) 
              output size torch.Size([4, 2])
    In Model: input size torch.Size([4, 5]) 
              output size torch.Size([4, 2])
    In Model: input size torch.Size([4, 5]) 
              output size torch.Size([4, 2])
    In Model: input size torch.Size([4, 5]) 
              output size torch.Size([4, 2])
Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
    In Model: input size torch.Size([4, 5]) 
              output size torch.Size([4, 2])
    In Model: input size torch.Size([4, 5]) 
              output size torch.Size([4, 2])
    In Model: input size torch.Size([4, 5]) 
              output size torch.Size([4, 2])
    In Model: input size torch.Size([4, 5]) 
              output size torch.Size([4, 2])
    In Model: input size torch.Size([4, 5]) 
              output size torch.Size([4, 2])
    In Model: input size torch.Size([4, 5]) 
              output size torch.Size([4, 2])
    In Model: input size torch.Size([2, 5]) 
              output size torch.Size([2, 2])
    In Model: input size torch.Size([4, 5]) 
              output size torch.Size([4, 2])
Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
    In Model: input size torch.Size([4, 5]) 
              output size torch.Size([4, 2])
    In Model: input size torch.Size([4, 5]) 
              output size torch.Size([4, 2])
    In Model: input size torch.Size([4, 5]) 
              output size torch.Size([4, 2])
    In Model: input size torch.Size([4, 5]) 
              output size torch.Size([4, 2])
    In Model: input size torch.Size([4, 5]) 
              output size torch.Size([4, 2])
    In Model: input size torch.Size([4, 5]) 
              output size torch.Size([4, 2])
    In Model: input size torch.Size([4, 5]) 
              output size torch.Size([4, 2])
    In Model: input size torch.Size([2, 5]) 
              output size torch.Size([2, 2])
Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
    In Model: input size torch.Size([2, 5]) 
              output size torch.Size([2, 2])
    In Model: input size torch.Size([2, 5]) 
              output size torch.Size([2, 2])
    In Model: input size torch.Size([2, 5]) 
              output size torch.Size([2, 2])
    In Model: input size torch.Size([2, 5]) 
              output size torch.Size([2, 2])
    In Model: input size torch.Size([2, 5]) 
              output size torch.Size([2, 2])
Outside: input size torch.Size([10, 5]) output_size torch.Size([10, 2])


Заключение

DataParallel разделяет ваши данные автоматически и отсылает рабочие задания на несколько моделей на нескольких GPU. После того как все модели завершат своя работу, DataParallel собирает и объединяет результаты перед тем как их вернуть.


Читайте также другие статьи в блоге:

PyTorch: тензоры (tensors)

PyTorch: нейронные сети

PyTorch: создание классификатора изображений