Класс DataParallel

Класс DataParallel из модуля torch.nn реализует параллельную обработку данных на нескольких GPU. При передаче тензора на вход такой модели он автоматически разбивается по нулевой размерности (batch dimension) на части, каждая из которых отправляется на отдельный GPU. Затем результаты работы модели собираются обратно в единый выходной тензор. Это позволяет увеличить эффективный размер пакета и ускорить обучение. Класс принимает на вход саму модель и список идентификаторов устройств, на которых будет производиться вычисление.

Синтаксис

torch.nn.DataParallel(
	module,
	device_ids=None,
	output_device=None,
	dim=0
)

Параметры

Основные параметры конструктора:

module - экземпляр модели (сеть), которую нужно распараллелить.
device_ids - список идентификаторов устройств (GPU), например [0, 1, 2]. По умолчанию используются все доступные устройства.
output_device - устройство, на которое будет собираться результат. По умолчанию используется первое устройство из device_ids.
dim - размерность, по которой разбивается входной тензор. По умолчанию dim=0 (размерность пакета).

Пример

Создадим простую линейную модель и обернём её в DataParallel для использования на двух GPU:

import torch
import torch.nn as nn

model = nn.Linear(10, 5)
if torch.cuda.device_count() >= 2:
	model = nn.DataParallel(model, device_ids=[0, 1])

print(model)

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

DataParallel(
	(module): Linear(in_features=10, out_features=5, bias=True)
)

Пример

Передадим данные в распараллеленную модель. Входной тензор будет автоматически разделён по первому измерению:

import torch
import torch.nn as nn

torch.manual_seed(0)
model = nn.Linear(10, 5)
if torch.cuda.device_count() >= 2:
	model = nn.DataParallel(model)

x = torch.randn(8, 10)
if torch.cuda.is_available():
	x = x.cuda()
	model = model.cuda()

res = model(x)
print(res.shape)

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

torch.Size([8, 5])

Пример

Используем модель с несколькими входами, обёрнутую в DataParallel. В этом случае все входные тензоры будут разделены по одинаковой размерности:

import torch
import torch.nn as nn

class MyModel(nn.Module):
	def __init__(self):
		super().__init__()
		self.fc = nn.Linear(10, 5)
	
	def forward(self, x, y):
		return self.fc(x) + y

torch.manual_seed(0)
model = MyModel()
if torch.cuda.device_count() >= 2:
	model = nn.DataParallel(model)

x = torch.randn(8, 10)
y = torch.randn(8, 5)
if torch.cuda.is_available():
	x = x.cuda()
	y = y.cuda()
	model = model.cuda()

res = model(x, y)
print(res.shape)

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

torch.Size([8, 5])

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

класс DistributedDataParallel,
который обеспечивает более эффективное распараллеливание для распределённого обучения
метод forward,
который выполняется внутри параллельной обёртки
функцию torch.cuda.device_count,
которая возвращает количество доступных GPU
функцию to,
которая переносит модель или тензор на указанное устройство