Pytorch. Как работает pin_memory в Dataloader? - PullRequest
Новая
       61

Pytorch. Как работает pin_memory в Dataloader?

2 голосов
/ 07 апреля 2019

Я хочу понять, как работает pin_memory в Dataloader.

Согласно документации. 

pin_memory (bool, optional) – If True, the data loader will copy tensors into CUDA pinned memory before returning them.

Ниже приведен пример автономного кода. 

import torchvision
import torch

print('torch.cuda.is_available()', torch.cuda.is_available())
train_dataset = torchvision.datasets.CIFAR10(root='cifar10_pytorch', download=True, transform=torchvision.transforms.ToTensor())
train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=64, pin_memory=True)
x, y = next(iter(train_dataloader))
print('x.device', x.device)
print('y.device', y.device)

Я получил вывод. 

torch.cuda.is_available() True
x.device cpu
y.device cpu

Но я ожидал чего-то подобного, потому что я указал флаг pin_memory=True в Dataloader. 

torch.cuda.is_available() True
x.device cuda:0
y.device cuda:0

Также я запускаю какой-то тест: 

import torchvision
import torch
import time
import numpy as np

pin_memory=True
train_dataset =torchvision.datasets.CIFAR10(root='cifar10_pytorch', download=True, transform=torchvision.transforms.ToTensor())
train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=64, pin_memory=pin_memory)
print('pin_memory:', pin_memory)
times = []
n_runs = 10
for i in range(n_runs):
    st = time.time()
    for bx, by in train_dataloader:
        bx, by = bx.cuda(), by.cuda()
    times.append(time.time() - st)
print('average time:', np.mean(times))

Я получил следующие результаты. 

pin_memory: False
average time: 6.5701503753662

pin_memory: True
average time: 7.0254474401474

Так что pin_memory=True только делает вещи медленнее. Может кто-нибудь объяснить мне это поведение?

1 Ответ

6 голосов
/ 08 апреля 2019

Документация, возможно, слишком лаконична, учитывая, что используемые термины довольно нишевые. В терминах CUDA закрепленная память означает не память графического процессора, а память невыгружаемого процессора. Преимущества и обоснование предоставляются здесь , но суть его в том, что этот флаг позволяет операции x.cuda() (которую вам все равно придется выполнять как обычно), чтобы избежать одной неявной копии ЦП, что делает его немного более производительным. Кроме того, с закрепленными тензорами памяти вы можете использовать x.cuda(non_blocking=True) для асинхронного копирования по отношению к хосту. Это может привести к повышению производительности в определенных сценариях, а именно, если ваш код структурирован как

  1. x.cuda(non_blocking=True)
  2. выполняет некоторые операции с процессором
  3. выполняет операции с графическим процессором, используя x.

Поскольку копия, инициированная в 1., является асинхронной, она не блокирует 2. от продолжения, пока идет копирование, и, таким образом, оба могут происходить рядом (что является усилением). Поскольку для шага 3. требуется, чтобы x уже был скопирован в графический процессор, его невозможно выполнить до тех пор, пока не будет завершено 1., поэтому только 1. и 2. могут перекрываться, а впоследствии 3. обязательно будет иметь место , Поэтому длительность 2. - это максимальное время, которое вы можете сэкономить с помощью non_blocking=True. Без non_blocking=True ваш ЦП будет бездействовать для завершения передачи, прежде чем продолжить с 2..

Примечание: возможно, шаг 2. может также включать операции с графическим процессором, если они не требуют x - я не уверен, верно ли это, и, пожалуйста, не цитируйте меня об этом.

Редактировать : Я полагаю, что вы упускаете точку в своем тесте Есть три проблемы с этим

  1. Вы не используете non_blocking=True в своих .cuda() вызовах.
  2. Вы не используете многопроцессорность в своем DataLoader, что означает, что большая часть работы выполняется в основном потоке в любом случае синхронно, что увеличивает затраты на передачу памяти.
  3. Вы не выполняете какую-либо работу ЦП в цикле загрузки данных (кроме вызовов .cuda()), поэтому нет необходимости накладываться на операции передачи памяти.

Эталоном, близким к тому, как предполагается использовать pin_memory, будет 

import torchvision, torch, time
import numpy as np

pin_memory = True
batch_size = 1024 # bigger memory transfers to make their cost more noticable
n_workers = 6 # parallel workers to free up the main thread and reduce data decoding overhead
train_dataset =torchvision.datasets.CIFAR10(
    root='cifar10_pytorch',
    download=True,
    transform=torchvision.transforms.ToTensor()
)   
train_dataloader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=batch_size,
    pin_memory=pin_memory,
    num_workers=n_workers
)   
print('pin_memory:', pin_memory)
times = []
n_runs = 10

def work():
    # emulates the CPU work done
    time.sleep(0.1)

for i in range(n_runs):
    st = time.time()
    for bx, by in train_dataloader:
       bx, by = bx.cuda(non_blocking=pin_memory), by.cuda(non_blocking=pin_memory)
       work()
   times.append(time.time() - st)
print('average time:', np.mean(times))

, что дает в среднем 5,48 с для моей машины с закреплением памяти и 5,72 с без.

...