The 2025 Dockerfile: PyTorch 2, CUDA 12 & NumPy 2

Let's get our hands dirty. A great Dockerfile is about making smart choices at every layer. We want an image that is reasonably sized, contains all the necessary tools, and builds reliably.

Choosing the Right Base Image

Everything starts with the FROM instruction. We'll use an official NVIDIA CUDA image. The key is choosing the right tag. You'll see -base, -runtime, and -devel. For a development environment, -devel is our best bet.

FROM nvidia/cuda:12.1.1-devel-ubuntu22.04

Why -devel? It includes the full CUDA toolkit, compilers (like NVCC), and debugging libraries. This is crucial if you need to build a Python package from source that has custom CUDA extensions—a common scenario in advanced ML research.

Environment Setup and Dependencies

Next, we'll set up some environment variables and install our base system packages using apt-get. We use DEBIAN_FRONTEND=noninteractive to prevent the build from hanging on any interactive prompts.

# Set a non-interactive frontend for package installations
ENV DEBIAN_FRONTEND=noninteractive

# Install essential system packages and Python
RUN apt-get update && apt-get install -y \
    python3.10 python3-pip git wget \
    && rm -rf /var/lib/apt/lists/*

Notice we chain the commands with && and clean up the apt cache with rm -rf /var/lib/apt/lists/* in the same RUN layer. This is a best practice for keeping image sizes down.

The Complete Dockerfile for 2025

Now, let's put it all together. This Dockerfile is our blueprint for a perfect ML development environment. The magic happens in the pip install step, where we carefully specify the versions and the PyTorch index URL to ensure we get a build that's compiled for CUDA 12.1.

# Base Image: Ubuntu 22.04 with CUDA 12.1 development toolkit
FROM nvidia/cuda:12.1.1-devel-ubuntu22.04

# --- Environment Setup ---

# Avoid prompts from apt during build
ENV DEBIAN_FRONTEND=noninteractive

# Set the working directory in the container
WORKDIR /app

# --- System Dependencies ---

# Install Python, pip, and other essential tools
RUN apt-get update && apt-get install -y \
    python3.10 \
    python3-pip \
    git \
    wget \
    # Clean up apt cache to reduce image size
    && rm -rf /var/lib/apt/lists/*

# --- Python Dependencies ---

# Install PyTorch, NumPy 2, and other core libraries
# Using --no-cache-dir reduces layer size
# Specifying the index URL for PyTorch is crucial for getting the correct CUDA-enabled build
RUN pip install --no-cache-dir --upgrade pip && \
    pip install --no-cache-dir \
    # Install PyTorch for CUDA 12.1
    torch==2.3.0 torchvision==0.18.0 torchaudio==2.3.0 --index-url https://download.pytorch.org/whl/cu121 && \
    # Install NumPy 2.0
    "numpy>=2.0.0" && \
    # Install other common ML libraries
    pip install --no-cache-dir \
    jupyterlab \
    scikit-learn \
    pandas \
    matplotlib \
    tqdm

# --- Final Setup ---

# Copy local code into the container
COPY . /app/

# Expose JupyterLab port
EXPOSE 8888

# Default command to start a bash shell
CMD ["bash"]

Build and Run Your New Environment

With the Dockerfile saved in your project root, building and running it is straightforward.

1. Build the image:

Open your terminal and run the build command. We'll tag it as ml-env-2025.

docker build -t ml-env-2025 .

2. Run the container:

Now, let's launch an interactive session inside our new environment. The --gpus all flag is essential for giving the container access to your NVIDIA GPUs.

docker run --gpus all -it --rm ml-env-2025 bash

3. Verify your setup:

Once you're inside the container's bash prompt, start a Python interpreter and run this snippet to confirm everything is working as expected:

import torch
import numpy as np

print(f"PyTorch Version: {torch.__version__}")
print(f"NumPy Version: {np.__version__}")
print(f"CUDA Available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"CUDA Version: {torch.version.cuda}")
    print(f"Current GPU: {torch.cuda.get_device_name(torch.cuda.current_device())}")

You should see output confirming your PyTorch and NumPy versions, and most importantly, CUDA Available: True. Success!

Final Thoughts: Reproducibility is Freedom

Building a robust Dockerfile isn't just a chore; it's an investment in your own productivity and sanity. This setup gives you a powerful, portable, and reproducible environment that you can use across all your projects, share with colleagues, and deploy to the cloud with confidence.

This Dockerfile is a fantastic starting point. Feel free to customize it by adding your own libraries to a requirements.txt file or modifying the base image for different needs. The era of "it works on my machine" is over. Welcome to the clean, containerized future of machine learning.