How to Set Up Docker with NVIDIA GPU for Deep Learning on Linux

Running deep‑learning models on a Linux server with a GPU is common, and using Docker for the same purpose requires extra steps so the container can see and use the host’s graphics card. This article explains how to set up the Docker GPU environment.

Nvidia Driver

First install the NVIDIA driver on the host. The easiest way is to download the installer script from the NVIDIA website, which works for any Linux distribution (CentOS, Ubuntu, etc.). The driver installation compiles a kernel module, so you need gcc and the kernel‑development packages, e.g., kernel-devel-$(uname -r).

Install gcc and kernel‑dev

$ sudo apt install gcc kernel-dev -y

Install NVIDIA driver

Visit NVIDIA driver download page .

Select the appropriate operating system and driver version, then click SEARCH.

Download the installer script and run it on the host.

$ wget https://www.nvidia.com/content/DriverDownload-March2009/confirmation.php?url=/tesla/450.80.02/NVIDIA-Linux-x86_64-450.80.02.run&lang=us&type=Tesla
$ chmod +x NVIDIA-Linux-x86_64-450.80.02.run && ./NVIDIA-Linux-x86_64-450.80.02.run

Verify the installation with nvidia-smi. If the output resembles the screenshot, the driver is installed correctly.

CUDA Driver

CUDA (Compute Unified Device Architecture) is NVIDIA’s parallel computing platform. Install CUDA similarly to the driver.

Download the appropriate CUDA toolkit from NVIDIA CUDA Toolkit Archive .

Configure the environment variable:

$ echo 'export PATH=/usr/local/cuda/bin:$PATH' | sudo tee /etc/profile.d/cuda.sh
$ source /etc/profile

nvidia-docker2

Docker installation is omitted here; refer to the official Docker documentation. This section focuses on installing nvidia-docker2.

nvidia-docker2 replaces the deprecated nvidia-docker1. Its components are:

libnvidia-container

nvidia-container-toolkit

nvidia-container-runtime

The runtime adds a hook that, when the NVIDIA_VISIBLE_DEVICES environment variable is set, mounts the GPU devices and CUDA driver into the container.

Installation Steps

Set the repository and GPG key:

$ distribution=$( . /etc/os-release; echo $ID$VERSION_ID )
$ curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -
$ curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list

Install the package and restart Docker:

$ sudo apt-get update
$ sudo apt-get install -y nvidia-docker2
$ sudo systemctl restart docker

Validate the setup:

Run the following command:

$ docker run --rm --gpus all nvidia/cuda:10.2-base nvidia-smi

If the output matches the host’s nvidia-smi result, the GPU is accessible inside containers.

For TensorFlow inside the container:

import tensorflow as tf
tf.contrib.eager.num_gpus()

For PyTorch inside the container:

import torch
torch.cuda.is_available()

Both commands should report the number of GPUs (or True for PyTorch), confirming GPU acceleration.

Usage Examples

Use all GPUs:

$ docker run --rm --gpus all nvidia/cuda nvidia-smi
$ docker run --rm --runtime=nvidia -e NVIDIA_VISIBLE_DEVICES=all nvidia/cuda nvidia-smi

Specify particular GPUs:

$ docker run --gpus "device=1,2" nvidia/cuda nvidia-smi
$ docker run --rm --runtime=nvidia -e NVIDIA_VISIBLE_DEVICES=1,2 nvidia/cuda nvidia-smi

With these steps, the basic environment for GPU‑accelerated computing in Docker is ready. Future work can explore GPU scheduling in Kubernetes.

This article is reproduced from “lxkaka” (original URL: https://lxkaka.wang/docker-nvidia/). All rights reserved by the original author.