Linux（ubuntu） GPU CUDA 构建Docker镜像

一、创建Dockerfile

FROM ubuntu:20.04

#非交互式，以快速运行自动化任务或脚本，无需图形界面
ENV DEBIAN_FRONTEND=noninteractive

# 安装基础工具
RUN apt-get update && apt-get install -y \
    curl \
    wget \
    git \
    build-essential \
    software-properties-common \
        python3.8 \
    python3-pip \
    && apt-get clean \
    && rm -rf /var/lib/apt/lists/*

RUN apt-get update && apt-get install -y python3-pip
# 更新 pip 和 setuptools
#RUN python3 -m pip install --upgrade pip setuptools

# 添加 NVIDIA CUDA 密钥和仓库
RUN apt-key adv --keyserver keyserver.ubuntu.com --recv-keys A4B469963BF863CC \
    && echo "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/ /" > /etc/apt/sources.list.d/cuda.list

# 安装 CUDA 和 cuDNN
RUN apt-get update && apt-get install -y --no-install-recommends \
    cuda-toolkit-12-4 \
    libcudnn8 \
    && apt-get clean \
    && rm -rf /var/lib/apt/lists/*

# 设置 Python 3.8 为默认 Python 版本
RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.8 1

# 安装gpu torch
RUN pip install --timeout=100000 torch==2.4.1 torchvision==0.19.1 torchaudio==2.4.1 \
    --index-url https://download.pytorch.org/whl/cu124 \
    -i https://mirrors.aliyun.com/pypi/simple

注意：cuda-toolkit以及torch相关的版本与cuda的版本都要相匹配，否则会安装失败。有其他需求则可自行更改Dockerfile文件。

二、启动容器并验证

启动容器：docker run -it --name 你的容器名称 --gpus all  你的镜像ID /bin/bash

查看显卡信息：nvidia-smi

验证GPU是否可用：python3 -c "import torch; print(torch.cuda.is_available())"