Настройка среды Conda tenorflow-GPU с CUDA 11.2 и cuDNN 8.1-8.2 (драйвер CUDA 460

У меня система Linux Mint 20.1 (основанная на Ubuntu 20.04 LTS) с GeForce RTX 3080 (версия драйвера 460.80), и у меня было много проблем при попытке запустить Tensorflow в среде conda.

Похоже, проблема в том, что в настоящее время не существует среды conda, которая правильно упакована с Tensorflow 2.4+, CUDA 11+ и CuDNN 8+, которые необходимы для работы на этой новой архитектуре GPU (более подробная информация здесь). Если вы используете conda install -c anaconda tensorflow-gpu, то по умолчанию будут установлены TF v2.2, cudatoolkit 10.x и cudnn 7.x. Если вы попытаетесь принудительно установить более новую версию, используя tensorflow-gpu=2.4, он либо просто установит старые несовместимые библиотеки cudatoolkit 10.x/cudnn 7.x, либо не установит их вообще.

Возможно, существует множество различных способов сделать это, но вот что сработало для меня после долгих проб и ошибок:

Шаг 1: Создайте среду conda и установите в нее cudatoolkit и cudnn.

conda create -n tf_gpu_env -c conda-forge cudatoolkit cudnn python=3.8

На момент написания статьи это позволит установить cudatoolkit 11.2, cudnn 8.2 и python 3.8.10 в эту новую среду. Я использовал канал conda-forge, но думаю, что каналы anaconda и nvidia тоже подойдут.

Шаг 2: Активируйте среду и установите tensorflow-gpu, используя pip не conda. Установка из conda либо займет очень много времени, так как conda пытается разрешить конфликты до выхода из строя, либо принудительно понизит cudatoolkit и cudnn до старых версий.

conda activate tf_gpu_env​
pip install tensorflow-gpu

На момент написания этой статьи здесь установлен Tensorflow-gpu 2.5.0

Шаг 3: Проверьте, что Tensorflow работает и использует GPU. Убедитесь, что у вас активировано новое окружение, и запустите сеанс python в терминале. Для проверки я использую следующие утверждения.

>>> import tensorflow as tf 

должен вернуть сообщение об успешном открытии libcudart

>>> tf.config.list_physical_devices('GPU')

должен вернуть длинное сообщение об успешном открытии кучи библиотек cuda и, что более важно, список в конце с именованным кортежем, указывающим, что он нашел GPU (например, [PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]). Если возвращается пустой список, значит, Tensorflow не использует GPU.

Наконец, создайте какой-нибудь случайный тензор с помощью tf.constant или tf.random. Эта проверка очень важна. Tensorflow может распознать ваш GPU даже если библиотеки cuda несовместимы и вернуть аналогичные сообщения при вводе двух вышеуказанных команд, указывая, что все в порядке. Однако, если все работает правильно, следующая команда (или аналогичная) должна выполняться и возвращать тензор почти мгновенно:

>>> tf.random.uniform([4, 4, 4, 4])

Если все не в порядке, будет очень большая задержка, прежде чем вы получите ответ (хотя последующие вызовы могут быть быстрыми). Это будет еще хуже, если вы попытаетесь запустить реальные модели с задержками, длящимися много минут или почти час до запуска первой эпохи, в дополнение к непредсказуемому поведению, такому как получение значений nan для некоторых сетей, таких как CNN.

Наконец, пара замечаний:

• Будьте осторожны при выполнении conda install или conda update в этой среде и тщательно проверяйте план пакета перед тем, как нажать клавишу Enter. Например, если вы используете conda для установки tensorflow-probability, он также может установить tensorflow-base в качестве зависимости, которая может перекрыть tensorflow-gpu.
• Вы также можете установить другие версии Tensorflow и библиотек cuda. Например, вы можете использовать TF 2.4.1, cudatoolkit 11.0 и cudnn 8.0, используя cudatoolkit=11.0, cudnn=8.0, tensorflow-gpu==2.4.1 (двойное равенство для pip) в командах установки выше.