Вот трюк для оптимизации нейросетки, который даёт примерно 4x ускорение при передаче данных с CPU на GPU.

Представим задачу классификации изображений.

Определяем модель, загружаем и преобразуем данные.

В цикле обучения передаём данные на GPU и тренируем сеть.

В чём проблема:

Если заглянуть в профайлер,

- большая часть ресурсов уходит на ядро (то есть на само обучение),
- но заметное время тратится и на передачу данных с CPU на GPU (cudaMemcpyAsync).

Сократить это просто.

Изначально датасет состоит из пикселей в виде 8-битных целых чисел. Мы преобразуем их в 32-битные float’ы.
А потом именно эти float-тензоры гоняем на GPU. В итоге данных стало в 4 раза больше, а значит, и передача тяжелее.

Решение:

Сдвинуть шаг преобразования после передачи. То есть сначала передаём 8-битные int’ы, а уже на GPU конвертируем их в float.

В результате шаг передачи данных заметно ускорится.

Конечно, это работает не везде, например, в NLP мы изначально имеем дело с float-эмбеддингами.
Но в случаях, где это применимо, ускорение получается очень ощутимым.

👉 @PythonPortal