Docker | PyTorch

When setting up a Docker image for PyTorch that supports both CPU and GPU, you'll need to include the necessary dependencies and configurations for both processing types. Here’s a basic setup you can use to create a Docker image that supports PyTorch on both CPU and GPU:

1. Base Image: Start with an official PyTorch image from the NVIDIA container repository, which comes with CUDA and cuDNN pre-installed, ensuring GPU support.

2. Requirements: Install any additional packages you need for your specific project.

3. Environment: Set up the correct environment variables and ensure that the container can dynamically switch between CPU and GPU based on the available hardware.

Here is a simple example of a Dockerfile:

# Use an official PyTorch image as a parent image
FROM pytorch/pytorch:1.12.0-cuda11.3-cudnn8-runtime

# Install any needed packages specified in requirements.txt
COPY requirements.txt /tmp/
RUN pip install --requirement /tmp/requirements.txt

# Run the application
CMD ["python", "./src/my_script.py"]

Key Notes:

• PyTorch Version: Replace 1.12.0-cuda11.3-cudnn8-runtime with the version of PyTorch that matches your requirements.
• CUDA/cuDNN: The CUDA and cuDNN versions need to match the ones required by the specific PyTorch version.
• requirements.txt: This file should contain all the Python packages you need. Copy this into your Docker image and install them.

To build and run the Docker container:

• Build the Image: docker build -t my-pytorch-app .
• Run the Container:
  • On CPU: docker run my-pytorch-app
  • On GPU: docker run --gpus all my-pytorch-app

This setup ensures that your Docker container can utilize GPU resources when available, falling back to CPU when GPUs are not accessibl

---

To set up a Docker image based on Ubuntu for PyTorch that supports both CPU and GPU usage, you'll start with an Ubuntu base image and manually install the necessary components like Python, PyTorch, CUDA, and cuDNN. This approach gives you more control over the versions and the setup.

Here’s a detailed Dockerfile to achieve this:

# Use an Ubuntu base image
FROM ubuntu:20.04

# Avoid prompts during package installation
ARG DEBIAN_FRONTEND=noninteractive

# Install basic utilities
RUN apt-get update && apt-get install -y \
    wget \
    curl \
    git \
    python3-pip \
    python3-dev \
    build-essential

# Link python3 to python
RUN ln -s /usr/bin/python3 /usr/bin/python

# Upgrade pip
RUN pip3 install --upgrade pip

# Install PyTorch with CUDA support
RUN pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113

# Copy the requirements file and install Python dependencies
COPY requirements.txt /tmp/
RUN pip3 install -r /tmp/requirements.txt

# Set the working directory
WORKDIR /app

# Copy the rest of your application's code
COPY . /app

# Command to run on container start
CMD ["python", "app.py"]

Explanation:

1. Base Image: Starts with Ubuntu 20.04.
2. Utilities: Installs essential packages like wget, curl, git, and Python-related tools.
3. Python Link: Makes python command available by linking python3 to python.
4. Python and PyTorch: Installs Python packages and PyTorch with CUDA 11.3 support.
5. Requirements: Installs additional Python dependencies from a requirements.txt file.
6. Working Directory and Application Code: Sets the working directory and copies your application code into the image.

Building and Running the Docker Container:

• Build the Image: docker build -t my-pytorch-app .
• Run the Container:
  • On CPU: docker run my-pytorch-app
  • On GPU: docker run --gpus all my-pytorch-app

This Dockerfile ensures that the application can utilize GPU acceleration when available, and will work on CPUs otherwise. Adjust the CUDA version in the PyTorch installation command according to the CUDA version compatible with your GPU drivers.