Питорч. Как работает pin_memory в Dataloader?

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

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

Поскольку копирование, инициированное в 1., является асинхронным, оно не блокирует выполнение 2. во время выполнения копирования, и, таким образом, оба могут происходить одновременно (что является преимуществом). Поскольку шаг 3. требует, чтобы x уже был скопирован в GPU, он не может быть выполнен, пока 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 с без.