How to Install PyTorch 2.4 with GPU Support on Docker

CUDA-capable GPU: Ensure you have an NVIDIA GPU installed.

Docker: Ensure that Docker is installed on your system. You can install Docker by following the official installation guide (https://docs.docker.com/engine/install/) for your operating system here.

NVIDIA Driver: Ensure that the appropriate NVIDIA driver is installed on your system. You can check the installation with the nvidia-smi command.

NVIDIA Container Toolkit: If you want GPU support in Docker containers, you need to install the NVIDIA Container Toolkit. This enables Docker to interface with NVIDIA GPUs. Installation Reference - How to Install NVIDIA Container Toolkit? ➥

PyTorch 2.4 is available on the PyTorch Docker Hub. Pull the specific image with CUDA support. These images will come with Python, PyTorch 2.4, and the necessary CUDA libraries pre-installed.

For CUDA 11.8:

$ docker pull pytorch/pytorch:2.4.0-cuda11.8-cudnn9-runtime
$ docker pull pytorch/pytorch:2.4.0-cuda11.8-cudnn9-devel

For CUDA 12.1:

$ docker pull pytorch/pytorch:2.4.0-cuda12.1-cudnn9-runtime
$ docker pull pytorch/pytorch:2.4.0-cuda12.1-cudnn9-devel

For CUDA 12.4:

$ docker pull pytorch/pytorch:2.4.0-cuda12.4-cudnn9-runtime
$ docker pull pytorch/pytorch:2.4.0-cuda12.4-cudnn9-devel

Note: What is the difference between devel and runtime tag for a Docker container?

runtime: extends the base image by adding all the shared libraries from the CUDA toolkit. Use this image if you have a pre-built application using multiple CUDA libraries. devel: extends the runtime image by adding the compiler toolchain, the debugging tools, the headers and the static libraries. Use this image to compile a CUDA application from sources.

Launch the Docker container with GPU support:

$ docker run --gpus all -it --rm pytorch/pytorch:2.4.0-cuda11.8-cudnn9-runtime

Inside the container, verify that PyTorch 2.4 is installed and that it can access the GPU:

import torch
print(torch.__version__)  # Should return 2.4.0
print(torch.cuda.is_available())  # Should return True

Compatibility: Ensure that your CUDA drivers are compatible with the CUDA version used in the PyTorch 2.4 Docker image.

Persistent Data: Use volume mounts to persist data outside the container as needed.

Pytorch provides several images in the Docker hub, which you can use by just pulling them. But in some cases, you may want to build a custom PyTorch Docker image that includes additional dependencies or configuration options.

If you need a custom Docker image with your specific environment or software, you can create a Dockerfile. Example Dockerfile:

FROM nvidia/cuda:12.4.0-base-ubuntu22.04

# Install necessary packages
RUN apt-get update && apt-get install -y \
    python3 \
    python3-pip

# Upgrade pip
RUN pip3 install --upgrade pip

# Install Python packages
RUN pip3 install numpy scipy

# Install any python packages you need
COPY path-in-host/requirements.txt requirements.txt
RUN pip3 install -r requirements.txt

# Install PyTorch and torchvision
RUN pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124

# Set the working directory
WORKDIR /home

# Set the command to run when the container starts
CMD ["nvidia-smi"]

This Dockerfile starts with the nvidia/cuda:12.4.0-base-ubuntu22.04 base image, which includes CUDA and cuDNN libraries. We then install system dependencies, including python3, numpy, scipy, and runs nvidia-smi by default. You can modify it to suit your needs.

You can build the Docker image by navigating to the directory containing the Dockerfile and running the following command:

# Create "pytorch-gpu" image from the Dockerfile
docker build -t pytorch-gpu . -f Dockerfile

The above command will build a Docker image named pytorch-gpu. The -t option is to allow you to give the image a name. The . means that the build context is the current directory and -f helps in pointing to the Dockerfile.

You have successfully built a Docker image, now you can run a container by executing the following command:

# Create and run a container from the above image  
docker run --name pytorch-container --gpus all -it --rm -v $(pwd):/home pytorch-gpu

The above command will start a Docker container named pytorch-container using the pytorch-gpu image. The --gpus all option specifies that all available GPUs should be exposed to the container. The -it option runs the container in interactive mode and attaches a terminal. The --rm option automatically removes the container when it exits. The -v option mounts the current directory as a volume inside the container at /home.

In this article, we discussed how to install PyTorch 2.4 on your GPU-enabled machine using Docker. You can use images provided on Docker Hub or customize your own Docker images. Both approaches allow you to create reproducible PyTorch environments that can be shared across different machines and platforms.