How to Deploy PyTorch on Docker Using Ubuntu 24.04

This guide is aimed at helping to get PyTorch up and running on Docker over Ubuntu 24.04. This tutorial is for absolute beginner and will walk you through verily compilation of installing and setting up Docker, pulling a PyTorch image, and running PyTorch container through out all the commands. At the end of this guide, you will have a PyTorch environment running on Docker, which is easily manageable and repeatable.

1. Introduction to Docker and PyTorch

What is Docker?

Docker provides an easy way to run, build and deploy applications in containers. What are Containers?Containers help to package the applications and all the dependencies in a standardised unit for development. As always, this guarantees that the application runs the same across environments, so it becomes a MUST tool for development and production

What is PyTorch?

PyTorch is an open-source machine learning library based on the Torch library. It is widely used for applications such as computer vision and natural language processing. PyTorch provides a flexible and efficient platform for building deep learning models, making it a popular choice among researchers and developers.

2. Installing Docker on Ubuntu 24.04

Before we can run PyTorch in Docker, we need to install Docker on our Ubuntu 24.04 system. Follow these steps to install Docker:

Step 1: Update Your System

First, update your system to ensure all packages are up to date. Open a terminal and run:

sudo apt update && apt upgrade

Step 2: Install Docker

Ubuntu 24.04 includes Docker in its official repositories. To install Docker, run the following commands:

sudo apt install docker.io

Step 3: Start and Enable Docker

After the installation, start the Docker service and enable it to start on boot:

sudo systemctl start docker
sudo systemctl enable docker

Step 4: Verify Docker Installation

To verify that Docker is installed correctly, run:

docker --version

You should see an output similar to Docker version 24.0.x, build xxxxx.

Step 5: Add User to Docker Group

To run Docker commands without sudo, add your user to the Docker group:

sudo usermod -aG docker $USER

Log out and log back in for the changes to take effect.

3. Setting Up a PyTorch Docker Container

With Docker installed, we can now set up a PyTorch Docker container.

Step 1: Pull the PyTorch Docker Image

Docker Hub hosts official PyTorch images. To pull the latest PyTorch image, use:

docker pull pytorch/pytorch:latest

This command will download the latest PyTorch image, which includes PyTorch, Python, and other dependencies.

Step 2: Run the PyTorch Container

To start a PyTorch container, use the following command:

docker run -it --name pytorch_container pytorch/pytorch:latest bash

This command will start a container named pytorch_container and open a bash shell inside it.

Step 3: Verify PyTorch Installation

Once inside the container, verify the PyTorch installation by running Python and importing PyTorch:

python3
import torch
print(torch.__version__)

You should see the PyTorch version number printed, confirming the installation.

4. Running and Managing PyTorch Containers

Step 1: Running Jupyter Notebook

The PyTorch Docker image comes with Jupyter Notebook installed. To run Jupyter Notebook inside the container, use the following command:

docker run -it -p 8888:8888 pytorch/pytorch:latest

This command will start a PyTorch container with Jupyter Notebook running and accessible at http://localhost:8888 on your host machine.

Step 2: Accessing Jupyter Notebook

Open your web browser and go to http://localhost:8888. You will be prompted to enter a token, which can be found in the terminal output when the container starts.

Step 3: Saving and Loading Notebooks

To save your work, you can mount a local directory to the container. This allows you to save Jupyter notebooks to your local file system. Use the -v option to mount a directory:

docker run -it -p 8888:8888 -v /path/to/local/dir:/workspace pytorch/pytorch:latest

Replace /path/to/local/dir with the path to the directory on your host machine where you want to save your notebooks.

Step 4: Stopping and Restarting Containers

To stop a running container, use the docker stop command followed by the container name or ID:

docker stop pytorch_container

To restart the container, use:

docker start pytorch_container

Step 5: Removing Containers

If you no longer need a container, you can remove it using the docker rm command:

docker rm pytorch_container

Step 6: Using GPU (Optional)

If you have a compatible NVIDIA GPU, you can leverage GPU acceleration for PyTorch by using the PyTorch Docker image with GPU support. First, ensure you have NVIDIA Docker runtime installed. Then, pull the PyTorch GPU image:

docker pull pytorch/pytorch:latest-gpu

Run the container with GPU support:

docker run --gpus all -it --name pytorch_gpu_container pytorch/pytorch:latest-gpu bash

Inside the container, you can verify PyTorch is using the GPU by running:

import torch
print("CUDA available: ", torch.cuda.is_available())

Step 7: Automating with Docker Compose

To manage your Docker containers more efficiently, consider using Docker Compose. Docker Compose allows you to define and run multi-container Docker applications. Here’s an example of a docker-compose.yml file for PyTorch:

version: '3.8'

services:
  pytorch:
    image: pytorch/pytorch:latest
    ports:
      - "8888:8888"
    volumes:
      - /home/user/my_notebooks:/workspace

To use this, save the file as docker-compose.yml and run:

docker-compose up

This command will start the PyTorch container with Jupyter Notebook and mount the specified volume.

Deploying PyTorch on Docker using Ubuntu 24.04 provides a flexible and efficient way to manage your machine learning environments. Docker ensures consistency across different setups and simplifies dependency management, making it an excellent tool for both development and production.

In this tutorial, we covered the basics of Docker, how to install Docker on Ubuntu 24.04, and how to set up and manage PyTorch containers. With this knowledge, you can now explore more advanced Docker features and customize your PyTorch environment to suit your specific needs.

Happy coding! If you have any questions or run into any issues, feel free to leave a comment below.