Функция активации Swish с большей эффективностью памяти — Стиль жизни

И как профилировать использование памяти графического процессора PyTorch

(Это переиздание этой записи из моего личного блога.)

Обновление на 2020-08-22: использование torch.cuda.max_memory_allocated() и torch.cuda.reset_peak_memory_stats() в более новой версии (1.6+) PyTorch, вероятно, более точно. «(ссылка)»

Мотивация

Недавно я пробовал модели EfficientNet, реализованные в PyTorch. Мне удалось успешно настроить предварительно обученные модели EfficientNet на моем наборе данных и достичь точности на уровне основных моделей, таких как SE-ResNeXt-50. Однако обучить модель с нуля оказалось намного сложнее.

Точно настроенные модели EfficientNet могут достигать той же точности при гораздо меньшем количестве параметров, но они, похоже, занимают много памяти графического процессора, чем следовало бы (по сравнению с обычными моделями). В репозитории Github есть открытая проблема, связанная с этой проблемой — [lukemelas / EfficientNet-PyTorch] Проблемы с памятью.

Пользователь Github @selina предположил, что пакетная нормализация и активация Swish являются узкими местами, и заявив, что, используя custom ops в PyTorch, мы можем уменьшить количество GPU. использование памяти до 30%.

Пользовательская функция Swish

Это наиболее простая реализация модуля активации Swish, используемого в EfficientNet (f (x) = x⋅σ (βx) с β = 1):

Градиенты этого модуля автоматически обрабатываются PyTorch.

Это модуль активации Swish, реализованный с помощью custom ops:

Здесь мы явно обрабатываем градиенты. Мы сохраняем копию входного тензора и используем ее на этапе обратного распространения для вычисления градиентов.

Действительно ли эта последняя версия значительно сокращает объем памяти графического процессора? — это вопрос, на который мы хотим ответить в следующем разделе.

Профилирование использования памяти CUDA

Я только что узнал, что теперь в PyTorch есть удобная функция torch.cuda.memory_allocated() , которую можно использовать для профилирования использования памяти графического процессора:

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

Примечание. Вероятно, это меньше, чем количество, указанное в nvidia-smi, поскольку некоторая неиспользуемая память может удерживаться кэширующим распределителем, а некоторый контекст необходимо создать на графическом процессоре. Дополнительные сведения об управлении памятью графического процессора см. В разделе «Управление памятью».

Я использовал для извлечения этой информации из вызовов команды nvidia-smi, но число, сообщаемое nvidia-smi, будет очень неточным из-за механизма кэширования памяти PyTorch.

Была создана простая полносвязная нейронная сеть для тестирования:

Мы просто вставили torch.cuda.memory_allocated() между операторами обучения модели, чтобы измерить использование памяти графического процессора. Для более сложного профилирования вам следует попробовать что-то вроде pytorch-memlab.

Наблюдения

При использовании пакетов размером 128, объем памяти графического процессора в цикле обучения был следующим:

(1st step) data: 524 MB forw: 1552 MB loss: 1552 MB back: 1044 MB step: 1044 MB ==================== (2nd step) data: 1044 MB forw: 2072 MB loss: 2072 MB back: 1044 MB step: 1044 MB (The latter steps are exactly the same as the second one.)

Разница между первой и последней эпохами, вероятно, связана с тем, что градиенты не выделяются до вызова loss.backward().

Пиковое использование памяти происходит сразу после прямого распространения. Как было показано в пользовательской реализации Swish, для некоторых функций требуется, чтобы PyTorch сохранял некоторые формы входных тензоров для возможности обратного распространения. Эта сохраненная информация удаляется после обратной фазы. По этой логике мы можем предположить, что обучение с большими размерами пакетов будет использовать больше памяти, и эта интуиция подтверждается экспериментами:

Пользовательская версия Swish использует почти на 20% меньше памяти, когда размер пакета составляет 512. PyTorch augograd, вероятно, решит сохранить больше информации в прямой фазе, чтобы избежать некоторых повторных вычислений в обратной фазе. Обратите внимание, что в версии с настраиваемой операцией i * (1 — sigmoid_i) в обратной функции можно рефакторировать, чтобы повторно использовать вычисленное число i * torch.sigmoid(i) в прямой функции.

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