Атрибут pin_memory

Атрибут pin_memory класса DataLoader управляет использованием закреплённой (page-locked) памяти для тензоров, возвращаемых загрузчиком. Если установить pin_memory=True, то каждый тензор, полученный из набора данных, будет скопирован в закреплённую память перед передачей в модель на GPU. Это позволяет ускорить операцию копирования данных с CPU на GPU благодаря использованию асинхронных операций и прямого доступа к памяти (DMA). Закреплённая память не может быть перемещена или заменена операционной системой, что делает её идеальной для частых и быстрых пересылок данных. Однако использование этой памяти увеличивает потребление системной оперативной памяти.

Синтаксис

torch.utils.data.DataLoader(
	dataset,
	pin_memory=False,  # по умолчанию выключен
	...
)

Атрибут задаётся в конструкторе загрузчика. Он принимает логическое значение (bool).

Пример с выключенным pin_memory

import torch
from torch.utils.data import DataLoader, TensorDataset

torch.manual_seed(0)
data = torch.randint(0, 10, (100, 3, 64, 64))
labels = torch.randint(0, 2, (100,))
dataset = TensorDataset(data, labels)
loader = DataLoader(dataset, batch_size=10, pin_memory=False)

for batch, labels in loader:
	print(batch.is_pinned())
	break

Результат выполнения кода:

False

Здесь метод is_pinned возвращает False, так как закрепление памяти отключено.

Пример с включенным pin_memory

import torch
from torch.utils.data import DataLoader, TensorDataset

torch.manual_seed(0)
data = torch.randint(0, 10, (100, 3, 64, 64))
labels = torch.randint(0, 2, (100,))
dataset = TensorDataset(data, labels)
loader = DataLoader(dataset, batch_size=10, pin_memory=True)

for batch, labels in loader:
	print(batch.is_pinned())
	break

Результат выполнения кода:

True

Теперь метод is_pinned возвращает True, что подтверждает закрепление памяти.

Использование с устройством CUDA

import torch
from torch.utils.data import DataLoader, TensorDataset

if torch.cuda.is_available():
	torch.manual_seed(0)
	data = torch.randint(0, 10, (100, 3, 64, 64))
	labels = torch.randint(0, 2, (100,))
	dataset = TensorDataset(data, labels)
	loader = DataLoader(
		dataset,
		batch_size=10,
		pin_memory=True,
		num_workers=2
	)
	
	device = torch.device('cuda')
	for batch, labels in loader:
		batch = batch.to(device, non_blocking=True)
		print(batch.device)
		break
else:
	print("CUDA недоступна")

Результат выполнения кода (при наличии GPU):

cuda:0

В этом примере загрузчик использует закреплённую память, а аргумент non_blocking=True в методе to позволяет выполнять асинхронное копирование на GPU, что ускоряет обучение.

Смотрите также

класс DataLoader,
который управляет загрузкой данных
атрибут num_workers,
который задаёт количество процессов для загрузки
атрибут batch_size,
который определяет размер батча
атрибут dataset,
который хранит ссылку на набор данных