13.6. Краткая реализация для нескольких GPU¶
Реализовывать параллелизм с нуля для каждой новой модели — занятие не из приятных. Более того, оптимизация инструментов синхронизации дает существенный выигрыш в производительности. В дальнейшем мы покажем, как это сделать с помощью высокоуровневых API фреймворков глубокого обучения. Математика и алгоритмы остаются теми же, что и в Section 13.5. Как и следовало ожидать, для запуска кода из этого раздела потребуется как минимум два GPU.
import torch
from torch import nn
from d2l import torch as d2l
13.6.1. Пример сети¶
Рассмотрим несколько более содержательную сеть, чем LeNet из Section 13.5, которая при этом остаётся достаточно простой и быстрой для обучения. Мы выберем вариант ResNet-18 (He et al., 2016). Поскольку входные изображения очень малы, мы внесём в архитектуру небольшие изменения. В частности, в отличие от Section 8.6, мы используем меньшее ядро свёртки, шаг и дополнение на начальном этапе. Кроме того, мы удаляем слой максимального пулинга.
#@save
def resnet18(num_classes, in_channels=1):
"""A slightly modified ResNet-18 model."""
def resnet_block(in_channels, out_channels, num_residuals,
first_block=False):
blk = []
for i in range(num_residuals):
if i == 0 and not first_block:
blk.append(d2l.Residual(out_channels, use_1x1conv=True,
strides=2))
else:
blk.append(d2l.Residual(out_channels))
return nn.Sequential(*blk)
# This model uses a smaller convolution kernel, stride, and padding and
# removes the max-pooling layer
net = nn.Sequential(
nn.Conv2d(in_channels, 64, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(64),
nn.ReLU())
net.add_module("resnet_block1", resnet_block(64, 64, 2, first_block=True))
net.add_module("resnet_block2", resnet_block(64, 128, 2))
net.add_module("resnet_block3", resnet_block(128, 256, 2))
net.add_module("resnet_block4", resnet_block(256, 512, 2))
net.add_module("global_avg_pool", nn.AdaptiveAvgPool2d((1,1)))
net.add_module("fc", nn.Sequential(nn.Flatten(),
nn.Linear(512, num_classes)))
return net
13.6.2. Инициализация сети¶
Мы будем инициализировать сеть внутри цикла обучения. Для освежения в памяти методов инициализации см. Section 5.4.
net = resnet18(10)
# Get a list of GPUs
devices = d2l.try_all_gpus()
# We will initialize the network inside the training loop
13.6.3. Обучение¶
Как и прежде, код обучения должен выполнять несколько базовых функций для эффективного параллелизма:
Параметры сети необходимо инициализировать на всех устройствах.
При итерации по набору данных мини-пакеты должны быть распределены между всеми устройствами.
Мы вычисляем функцию потерь и её градиент параллельно на устройствах.
Градиенты агрегируются, и параметры обновляются соответствующим образом.
В конце мы вычисляем точность (опять же параллельно), чтобы отчитаться об итоговой производительности сети. Процедура обучения во многом схожа с реализациями из предыдущих глав, за исключением того, что нам необходимо разделять и агрегировать данные.
def train(net, num_gpus, batch_size, lr):
train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)
devices = [d2l.try_gpu(i) for i in range(num_gpus)]
def init_weights(module):
if type(module) in [nn.Linear, nn.Conv2d]:
nn.init.normal_(module.weight, std=0.01)
net.apply(init_weights)
# Set the model on multiple GPUs
net = nn.DataParallel(net, device_ids=devices)
trainer = torch.optim.SGD(net.parameters(), lr)
loss = nn.CrossEntropyLoss()
timer, num_epochs = d2l.Timer(), 10
animator = d2l.Animator('epoch', 'test acc', xlim=[1, num_epochs])
for epoch in range(num_epochs):
net.train()
timer.start()
for X, y in train_iter:
trainer.zero_grad()
X, y = X.to(devices[0]), y.to(devices[0])
l = loss(net(X), y)
l.backward()
trainer.step()
timer.stop()
animator.add(epoch + 1, (d2l.evaluate_accuracy_gpu(net, test_iter),))
print(f'test acc: {animator.Y[0][-1]:.2f}, {timer.avg():.1f} sec/epoch '
f'on {str(devices)}')
Давайте посмотрим, как это работает на практике. В качестве разминки мы обучим сеть на одном GPU.
train(net, num_gpus=1, batch_size=256, lr=0.1)
Далее мы используем 2 GPU для обучения. По сравнению с LeNet, рассмотренной в Section 13.5, модель ResNet-18 значительно сложнее. Именно здесь проявляется преимущество параллелизма. Время вычислений существенно превышает время синхронизации параметров. Это улучшает масштабируемость, поскольку накладные расходы на параллелизм становятся менее значимыми.
train(net, num_gpus=2, batch_size=512, lr=0.2)
13.6.4. Итог¶
Данные автоматически оцениваются на устройствах, где они находятся.
Убедитесь, что сети инициализированы на каждом устройстве перед тем, как пытаться получить доступ к его параметрам. В противном случае возникнет ошибка.
Алгоритмы оптимизации автоматически агрегируют данные по нескольким GPU.
13.6.5. Упражнения¶
В этом разделе используется ResNet-18. Попробуйте изменить количество эпох, размер пакета и скорость обучения. Используйте больше графических процессоров для вычислений. Что произойдет, если попробовать это с 16 графическими процессорами (например, на экземпляре AWS p2.16xlarge)?
Иногда различные устройства обеспечивают разную вычислительную мощность. Мы можем одновременно использовать графические процессоры (GPU) и центральный процессор (CPU). Как следует распределить между ними задачи? Стоит ли это усилий? Почему да? Почему нет?