Установка драйверов Nvidia для Tensorflow и PyTorch для Ubuntu

Установка Tensorflow и PyTorch GPU для видеокарт Nvidia для Ubuntu

Чтобы ускорить обучения нейронных сетей их надо делать на видеокартах. CUDA - это технология от Nvidia, которая позволяет программировать вычисления на видеокартах Nvidia.

Воспользутесь инструкцией, если вам нужно установить драйвер Nvidia в LXC или докер контейнер.

Установка драйверов Nvidia

Убедитесь, что у вас установлена видеокарта Nvidia:

nvidia-smi

Для системы Ubuntu с видеокартой Nvidia нужно поставить драйвер. На момент написания данного поста последний драйвер 525.147.05.

Установите ядро:

aptitude update
aptitude install linux-generic linux-headers-generic linux-image-generic

Установите драйвер Nvidia:

aptitude install nvidia-driver-525 linux-modules-nvidia-525-generic

Для ядра hwe

aptitude install linux-generic-hwe-22.04 linux-headers-generic-hwe-22.04 linux-modules-nvidia-525-generic-hwe-22.04

Удалите старый ключ от Nvidia репозитория, если он у вас был подключен до этого:

apt-key del 7fa2af80

Подключите репозиторий Nvidia:

wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/cuda-ubuntu1804.pin -O "/etc/apt/preferences.d/cuda-repository-pin-600"
echo "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/ /" > /etc/apt/sources.list.d/nvidia.list

Установите новый ключ от Nvidia репозитория:

apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/3bf863cc.pub

Установите CUDA 11:

aptitude update
aptitude install cuda-cudart-11-7 libcublas-11-7 libcufft-11-7 libcufile-11-7 libcurand-11-7 libcusolver-11-7 libcusparse-11-7 libcudnn8 cuda-cufft-10-2

Можете также установить полный пакет CUDA:

aptitude install cuda-11-7

Установка PyTorch

Установка:

pip3 install torch torchvision torchaudio numpy matplotlib

Проверьте доступность GPU в PyTorch:

python3 -c "import torch; print(torch.cuda.is_available())"

Установка Tensorflow

Для начала, установите основные библиотеки Tensorflow GPU, Keras, Numpy, Matplotlib. В командой строке введите команды:

pip3 install tensorflow-gpu tensorflow_probability tensorflow_addons
pip3 install keras
pip3 install numpy
pip3 install matplotlib

Проверьте версию Tensorflow:

python3.8 -c "import tensorflow as tf; print(tf.__version__)"

Проверьте список драйверов:

python3.8 -c "import tensorflow as tf; print(tf.config.list_physical_devices('GPU'))"

Если без ошибок библиотека выдала список графических процессоров, значит все установилось корректно.

В начале вашей программы на python для Tensorflow должны быть следующие строки:

import os, sys

# Устанавливает переменные окружения
#os.environ["TF_GPU_ALLOCATOR"]="cuda_malloc_async"
os.environ["TF_CPP_VMODULE"]="gpu_process_state=10,gpu_cudamallocasync_allocator=10"

import tensorflow as tf

# Устанавливает размер памяти GPU в Мбайт
gpu_memory_limit = 1024
gpus = tf.config.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(gpus[0], True)
tf.config.experimental.set_virtual_device_configuration(
    gpus[0],
    [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=gpu_memory_limit)])
logical_gpus = tf.config.experimental.list_logical_devices('GPU')

И при обучении нейронной сети ставьте batch_size не более 128. При большом batch_size потребляется много видеопамяти, и выходят ошибки, о том что CUDA не может выделить память. CUDA_ERROR_OUT_OF_MEMORY

Установка Tensorflow на CPU

Если вы хотите установить нейронную сеть на CPU, то выполните команды:

pip3 uninstall tensorflow-gpu
pip3 install tensorflow-cpu tensorflow_probability tensorflow_addons

Это полезно, если вы используете Tensorflow JS, а пишите нейронную сеть на PyTorch.

Материалы

  1. Настройка среды для работы с Deep Learning в браузере, используя Облачную операционную систему
  2. Установка и настройка драйверов nvidia в LXC контейнере
  3. Список видеокарт Nvidia, которые поддерживают CUDA
  4. CUDA Linux Repository Key Rotation
  5. Установка cudnn