Использование GPU при работе с библиотекой

Для того чтобы использовать видеоадаптер (GPU) для работы с библиотекой (по умолчанию используется CPU) неоходимо выполнить следующие шаги:

  1. Перейти на сайт и выбрать необходимую версию PyTorch с поддержкой GPU для CUDA, которая установлена у вас в системе.

Выбор версии PyTorch с поддержкой GPU


Либо установить версию с помощью прямой команды (актуальная версия PyTorch с поддержкой CUDA 11.8):

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
  1. Убедитесь, что библиотека PyTorch с поддержкой GPU установилась корректно (все пакеты, относящиеся к torch должны иметь в версиях подпись +cu118 или версию, которую вы выбрали):

pip list
Установленные версии пакетов


  1. В случае, если пакеты установились некорректно, либо они уже были установлены, но без поддержки GPU, то необходимо удалить зависимости PyTorch, а затем установить их как указано в шаге №1:

pip uninstall pytorch, torchvision и torchaudio