Если вы регулярно используете 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, у нас много интересных статей, гайдов, моделей из мира машинного обучения.

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