2 окт 2024 в 20:36

Memory pinning для ускорения обучения моделей машинного обучения

Средний

Машинное обучение * Data Engineering * Python *

Перевод

Ожидает приглашения

Автор оригинала: https://x.com/_avichawla

Если вы регулярно используете GPU для обучения моделей, то существует довольно простая техника, которая часто используется для ускорения обучения моделей.

...изменив всего две строки кода.

Рассмотрим стандартный цикл обучения модели в PyTorch:

В этом коде:

Строка 5 передает данные в GPU из CPU.
Все выполняется на GPU после передачи данных, то есть в строках 7-15.

Это означает, что когда работает GPU, CPU простаивает, а когда работает CPU, GPU простаивает, наглядно:

Но вот что мы можем сделать, чтобы оптимизировать:

Когда модель обучается на 1-м батче, CPU может передать 2-й батч на GPU.
Таким образом, GPU не придется ждать следующего батча данных, как только он завершит обработку существующего батча.

Иными словами, график использования ресурсов будет выглядеть примерно так:

В то время, когда CPU будет простаивать, GPU (который является фактическим ускорителем для обучения модели) гарантированно будет иметь данные для работы.

Формально этот процесс известен как memory pinning, и он используется для ускорения передачи данных от CPU к GPU, делая процесс обучения асинхронным.

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

Включить эту функцию в PyTorch довольно просто. Во-первых, при определении объекта DataLoader надо установить pin_memory=True и указать num_workers:

Далее, на этапе передачи данных в шаге обучения укажите non_blocking=True:

Готово!

Вот как работает ускорение на примере набора данных MNIST в обучении простой нейронной сети:

Без memory pinning обучение модели на 5 эпохах занимает около 43 секунд:

а с использованием memory pinning та же модель обучается менее чем за 10 (!!!) секунд:

Важные особенности использования memory pinning:

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

Поэтому, всякий раз, когда используете memory pinning - отслеживайте потребление RAM!

когда набор данных относительно мал, memory pinning имеет незначительный эффект, поскольку передача данных от CPU к GPU все равно не занимает столько времени:

👉 Переведено на русский язык редакцией телеграм-канала Machine learning interview, у нас много интересных статей, гайдов, моделей из мира машинного обучения.

Спасибо за внимание.

Хабы: