При обучении на GPU данные переносятся на устройство через
.to(device), после чего все вычисления выполняются уже на видеокарте. В итоге процесс получается неэффективным: когда работает GPU, CPU простаивает, и наоборот.Использование фиксации памяти решает эту проблему. Пока модель обучается на первом мини-батче, CPU успевает заранее перенести второй мини-батч на GPU. Таким образом, видеокарта не простаивает в ожидании новых данных и сразу начинает вычисления с готовым батчем.
Реализовать это просто:
- добавить
pin_memory=True в DataLoader,- при передаче данных использовать
.to(device, non_blocking=True)Кроме того, рекомендуется задать
num_workers в DataLoader для параллельной загрузки данных.Тесты показывают, что такие настройки дают заметное ускорение, что подтверждается результатами на графиках.
👉 @PythonPortal