Упростите обучение смешанной точности с MXNet-AMP

В то время как обучение в FP16 показало большой успех в задачах классификации изображений, другие более сложные нейронные сети обычно оставались в FP32 из-за трудностей с ручным применением рекомендаций по обучению со смешанной точностью. Именно здесь в игру вступает автоматическая смешанная точность (AMP): она автоматически применяет рекомендации обучения смешанной точности, устанавливая соответствующий тип данных для каждого оператора.

В этом блоге мы покажем, как начать обучение со смешанной точностью с использованием AMP в MXNet, на примере сети SSD от GluonCV.

Внутри автоматической смешанной точности

AMP имеет 2 основные особенности:

• Автоматическое литье модели до смешанной точности
• Масштабирование динамических потерь с помощью предоставленного набора утилит

На этапе автоматического преобразования AMP ищет в модели операции, которые лучше всего подходят для выполнения с более низкой точностью (например, слои Conv или Dense), преобразовывая их в FP16 (показано зеленым цветом на рисунке [3] а). С другой стороны, AMP также определяет операции, которые должны оставаться с полной точностью (например, Norm или exp), и заставляет их оставаться в FP32 (показано красным на рис. [3] а). Затем он распространяет типы данных через граф вычислений (рисунок [3] b). Всякий раз, когда встречается оператор с несколькими входами (например, сложение или умножение), AMP приводит все эти входные данные к самому широкому из их типов (показанному как крайний правый красный узел графа на рисунке [3] c). Например, если оператор имеет входы FP16 и FP32, все они будут преобразованы в FP32.

Уменьшенный динамический диапазон [4] FP16 по сравнению с FP32 может привести к тому, что некоторые градиенты будут слишком большими (переполнение) или слишком маленькими (потеря значимости) для правильного представления. Обе эти ситуации могут быть смягчены с помощью масштабирования потерь - метода, который сдвигает диапазон градиентов вверх или вниз, чтобы удерживать их в динамическом диапазоне FP16 (рисунок [4]). Чтобы быть наиболее эффективным для широкого спектра различных обучающих задач и сетей глубокого обучения, коэффициент масштабирования потерь необходимо выбирать динамически, реагируя на изменение величины градиента на протяжении всего обучения. AMP предоставляет API для простой реализации этого динамического масштабирования потерь в обучающем сценарии.

AMP в MXNet

Типичный цикл обучения в MXNet Gluon выглядит так:

net = …
loss = …
trainer = mx.gluon.Trainer(…)
for data, label in data_iter:
    with mx.autograd.record():
        out = net(data)
        l = loss(out, label)
        mx.autograd.backward(l)
    trainer.step()

Чтобы использовать AMP, необходимо внести всего несколько изменений, которые выделены жирным шрифтом:

from mxnet.contrib import amp
amp.init()
net = …
loss = …
trainer = mx.gluon.Trainer(…)
amp.init_trainer(trainer)
for data, label in data_iter:
    with mx.autograd.record():
        out = net(data)
        l = loss(out, label)
        with amp.scale_loss(l, trainer) as scaled_loss:
            mx.autograd.backward(scaled_loss)
trainer.step()

Посмотрим, в свою очередь, что делают эти дополнительные строки. Во-первых, у нас есть:

from mxnet.contrib import amp
amp.init()

Эта часть инициализирует AMP и изменяет поведение операторов, чтобы реализовать преимущества смешанной точности. Далее идет:

amp.init_trainer(trainer)

Это инициализирует Gluon Trainer для использования с динамическим масштабированием потерь, встроенным в AMP. И наконец:

with amp.scale_loss(l, trainer) as scaled_loss:
    mx.autograd.backward(scaled_loss)

Эти 2 строки выполняют часть масштабирования динамических потерь рецепта смешанной точности и вычисляют градиенты относительно масштабированных потерь. Обратите внимание, что градиенты, полученные с использованием этой масштабированной версии потерь, также масштабируются на то же значение. Градиенты немасштабируются внутри функции trainer.step () во время фазы обновления веса. Однако такое поведение может быть нежелательным, когда обучающие сценарии выполняют манипуляции с градиентом, такие как отсечение, перед этапом оптимизатора. Для использования в этих случаях AMP предоставляет функцию amp.unscale (trainer), которая предписывает AMP выполнить масштабирование градиента раньше, прежде чем может произойти какое-либо подобное манипулирование градиентом.

AMP на практике

Детектор одиночного выстрела

В качестве конкретного примера мы рассмотрим, как ускорить обучение модели обнаружения SSD (Single Shot MultiBox Detector, Лю и др. 2015)

Для этого эксперимента мы будем использовать существующую модель и обучающий сценарий из Модельного зоопарка GluonCV. Используемый графический процессор представляет собой один графический процессор V100 объемом 32 ГБ, доступный, например, в экземплярах AWS P3DN.

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

python train_ssd.py — dataset coco -j 4 — gpus 0 — network resnet50_v1 — data-shape 512
INFO:root:[Epoch 0][Batch 49], Speed: 58.105 samples/sec, CrossEntropy=1.190, SmoothL1=0.688
INFO:root:[Epoch 0][Batch 99], Speed: 58.683 samples/sec, CrossEntropy=0.693, SmoothL1=0.536
INFO:root:[Epoch 0][Batch 149], Speed: 58.915 samples/sec, CrossEntropy=0.500, SmoothL1=0.453
INFO:root:[Epoch 0][Batch 199], Speed: 58.422 samples/sec, CrossEntropy=0.396, SmoothL1=0.399

Базовая производительность составляет около 58 выборок в секунду. Давайте теперь проведем тот же эксперимент с использованием AMP. Первый шаг - включить AMP и инициализировать его:

from mxnet.contrib import amp
amp.init()

Следующим шагом является создание трейнера и его инициализация с помощью AMP для динамического масштабирования потерь. В настоящее время поддержка динамического масштабирования потерь ограничена тренерами, созданными с параметром update_on_kvstore = False, поэтому мы добавляем его в инициализацию нашего трейнера:

trainer = gluon.Trainer(net.collect_params(), ‘sgd’,
    {‘learning_rate’: args.lr, ‘wd’: args.wd, ‘momentum’: args.momentum},
    update_on_kvstore=False))
amp.init_trainer(trainer)

Последняя модификация заключается в добавлении фактического масштабирования динамических потерь:

with amp.scale_loss(sum_loss, trainer) as scaled_loss:
    autograd.backward(scaled_loss)

Вот и все. Однако, если вы попытаетесь запустить сценарий напрямую, вы можете получить следующую ошибку:

mxnet.base.MXNetError: [21:49:01] src/operator/nn/convolution.cc:297: Check failed: (*in_type)[i] == dtype (0 vs. 2) : This layer requires uniform type. Expected ‘float16’ v.s. given ‘float32’ at ‘weight’

Эта ошибка возникает из-за того, что сценарий SSD от GluonCV перед фактическим обучением запускает сеть один раз в контексте ЦП, чтобы получить якоря для загрузчика данных, а контекст ЦП не поддерживает некоторые операции FP16, такие как Слои Conv или Dense. Мы исправим это, изменив функцию get_dataloader (), чтобы использовать контекст GPU для генерации привязки:

def get_dataloader(net, train_dataset, val_dataset, data_shape, batch_size, num_workers, ctx):
    […]
    net.collect_params().reset_ctx(ctx)
    with autograd.train_mode():
        _, _, anchors = net(mx.nd.zeros((1, 3, height, width), ctx))
        anchors = anchors.as_in_context(mx.cpu())

Теперь мы можем запустить этот модифицированный скрипт и наблюдать за увеличением скорости.

python amp_train_ssd.py — dataset coco -j 4 — gpus 0 — network resnet50_v1 — data-shape 512
INFO:root:[Epoch 0][Batch 49], Speed: 93.585 samples/sec, CrossEntropy=1.166, SmoothL1=0.684
INFO:root:[Epoch 0][Batch 99], Speed: 93.773 samples/sec, CrossEntropy=0.682, SmoothL1=0.533
INFO:root:[Epoch 0][Batch 149], Speed: 93.399 samples/sec, CrossEntropy=0.493, SmoothL1=0.451
INFO:root:[Epoch 0][Batch 199], Speed: 93.674 samples/sec, CrossEntropy=0.391, SmoothL1=0.397

Производительность с AMP увеличивается с 58 выборок в секунду до 93 выборок в секунду. Это показывает, что, используя MXNet-AMP и лишь с небольшими дополнениями кода, мы получаем увеличение скорости на 60% по сравнению со сценарием FP32 по умолчанию. Вы можете сравнить обучение с AMP и без него с параметром - amp в сценарии обучения.

Маска-RCNN

Мы также можем добавить AMP к другой задаче глубокого обучения: сегментации экземпляров. Самая популярная модель - Маска-RCNN He et al. 2017, который может обрабатывать 6 выборок в секунду в исходной реализации, как показано в журнале обучения ниже.

Давайте посмотрим на производительность до AMP на экземпляре AWS P3DN с одним графическим процессором V100 32 ГБ:

python train_mask_rcnn.py — dataset coco -j 4 — gpus 0
INFO:root:[Epoch 0][Batch 49], Speed: 4.749 samples/sec, RPN_Conf=0.402,RPN_SmoothL1=0.155,RCNN_CrossEntropy=10.085,RCNN_SmoothL1=2.458,RCNN_Mask=5.818,RPNAcc=0.863,RPNL1Loss=0.628,RCNNAcc=0.808,RCNNL1Loss=1.928,MaskAcc=0.547,MaskFGAcc=0.573
INFO:root:[Epoch 0][Batch 99], Speed: 5.951 samples/sec, RPN_Conf=0.335,RPN_SmoothL1=0.152,RCNN_CrossEntropy=8.767,RCNN_SmoothL1=2.508,RCNN_Mask=5.451,RPNAcc=0.877,RPNL1Loss=0.614,RCNNAcc=0.821,RCNNL1Loss=1.929,MaskAcc=0.588,MaskFGAcc=0.583
INFO:root:[Epoch 0][Batch 149], Speed: 5.916 samples/sec, RPN_Conf=0.310,RPN_SmoothL1=0.150,RCNN_CrossEntropy=8.174,RCNN_SmoothL1=2.534,RCNN_Mask=5.234,RPNAcc=0.881,RPNL1Loss=0.598,RCNNAcc=0.826,RCNNL1Loss=1.934,MaskAcc=0.615,MaskFGAcc=0.591
INFO:root:[Epoch 0][Batch 199], Speed: 5.825 samples/sec, RPN_Conf=0.296,RPN_SmoothL1=0.149,RCNN_CrossEntropy=7.885,RCNN_SmoothL1=2.555,RCNN_Mask=5.088,RPNAcc=0.884,RPNL1Loss=0.592,RCNNAcc=0.826,RCNNL1Loss=1.925,MaskAcc=0.633,MaskFGAcc=0.595

Еще раз, есть три простых модификации:

• Import AMP:
    from mxnet.contrib import amp
    amp.init()

• Инициализируйте трейнер с помощью AMP:

trainer = gluon.Trainer(net.collect_params(), ‘sgd’,
    {‘learning_rate’: args.lr, ‘wd’: args.wd, ‘momentum’: args.momentum},
    update_on_kvstore=False))
amp.init_trainer(trainer)

• Выполните масштабирование потерь:

with amp.scale_loss(losses, trainer) as scaled_losses:
    autograd.backward(scaled_losses)

Давайте посмотрим на скорость с AMP:

python amp_train_mask_rcnn.py — dataset coco -j 4 — gpus 0
INFO:root:[Epoch 0][Batch 49], Speed: 9.043 samples/sec, RPN_Conf=0.449,RPN_SmoothL1=0.148,RCNN_CrossEntropy=12.762,RCNN_SmoothL1=2.483,RCNN_Mask=5.921,RPNAcc=0.842,RPNL1Loss=0.598,RCNNAcc=0.725,RCNNL1Loss=1.952,MaskAcc=0.542,MaskFGAcc=0.568
INFO:root:[Epoch 0][Batch 99], Speed: 9.779 samples/sec, RPN_Conf=0.360,RPN_SmoothL1=0.147,RCNN_CrossEntropy=9.857,RCNN_SmoothL1=2.512,RCNN_Mask=5.487,RPNAcc=0.867,RPNL1Loss=0.595,RCNNAcc=0.781,RCNNL1Loss=1.945,MaskAcc=0.585,MaskFGAcc=0.582
INFO:root:[Epoch 0][Batch 149], Speed: 10.281 samples/sec, RPN_Conf=0.328,RPN_SmoothL1=0.145,RCNN_CrossEntropy=8.883,RCNN_SmoothL1=2.558,RCNN_Mask=5.245,RPNAcc=0.874,RPNL1Loss=0.578,RCNNAcc=0.798,RCNNL1Loss=1.945,MaskAcc=0.614,MaskFGAcc=0.586
INFO:root:[Epoch 0][Batch 199], Speed: 9.603 samples/sec, RPN_Conf=0.309,RPN_SmoothL1=0.143,RCNN_CrossEntropy=8.290,RCNN_SmoothL1=2.576,RCNN_Mask=5.068,RPNAcc=0.879,RPNL1Loss=0.572,RCNNAcc=0.806,RCNNL1Loss=1.927,MaskAcc=0.636,MaskFGAcc=0.586

Включая AMP, мы получаем около 10 выборок в секунду. Это дает нам увеличение скорости на 67% по сравнению с версией FP32 по умолчанию.

Контрольные точки

В наших тестах с моделями SSD и Faster RCNN AMP продемонстрировал примерно двукратное улучшение пропускной способности обучения (выборок в секунду) с одним экземпляром графического процессора. В наших предварительных тестах Multi-GPU с Horovod, AMP улучшила скорость обучения модели SSD примерно на 20%.

использованная литература

[1] Используйте AMP (автоматическую смешанную точность) в MXNet https://mxnet.apache.org/api/python/docs/tutorials/performance/backend/amp.html

[2] Автоматическая смешанная точность для глубокого обучения, https://developer.nvidia.com/automatic-mixed-precision

[3] Формат с плавающей запятой одинарной точности, Википедия, https://en.wikipedia.org/wiki/Single-precision_floating-point_format

[4] Обучение со смешанной точностью, https://docs.nvidia.com/deeplearning/sdk/mixed-precision-training/index.html