一、创建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())"