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

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

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

Рассмотрим стандартный цикл обучения модели в PyTorch (См Картинку 1)

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

💡 Это означает, что когда работает GPU, CPU простаивает, а когда работает CPU, GPU простаивает, наглядно(См Картинку 2)

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

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

👉 Иными словами, график использования ресурсов будет выглядеть примерно так (См Картинку 3)

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

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

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

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

Далее, на этапе передачи данных в шаге обучения укажите non_blocking=True (См Картинку 5)

⚡️ Готово!

Вот как работает ускорение на примере набора данных MNIST в обучении простой нейронной сети (См Картинку 6)

- Без memory pinning обучение модели на 5 эпохах занимает около 43 секунд:
- а с использованием memory pinning та же модель обучается менее чем за 10 (!!!) секунд 🔥(См Картинку 7)

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

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

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

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

📌 Полная версия

@ai_machinelearning_big_data