在Docker容器中用Python3.11安装PyTorch+CUDA遇问题求助

有的，很多开发者已经在Python 3.11环境下成功部署PyTorch+CUDA。你遇到的问题核心是直接使用pip install torch时默认安装了CPU版本的PyTorch，而非适配CUDA的版本——因为Python官方镜像本身不包含CUDA依赖环境，pip会自动选择兼容当前环境的CPU包。而pytorch/pytorch官方镜像预先配置了完整的CUDA依赖链，因此安装时会自动匹配CUDA版本的PyTorch。

解决方案1：修正Python 3.11 Docker镜像的安装步骤 #

修改你的Dockerfile，明确指定安装适配CUDA版本的PyTorch包，同时补充必要的CUDA依赖库：

FROM python:3.11.4
# 安装CUDA运行时依赖（以CUDA 12.1为例，可根据宿主机版本调整）
RUN apt-get update && apt-get install -y --no-install-recommends \
    libcuda1-12.1 \
    libcudnn8=8.9.2.26-1+cuda12.1 \
    && rm -rf /var/lib/apt/lists/*
# 指定CUDA版本安装PyTorch
RUN --mount=type=cache,id=pip-build,target=/root/.cache/pip \
    pip install torch torchaudio --index-url https://download.pytorch.org/whl/cu121
ENV PATH="/usr/local/nvidia/bin:${PATH}" \
    NVIDIA_VISIBLE_DEVICES=all \
    NVIDIA_DRIVER_CAPABILITIES=all

• 关键说明：--index-url https://download.pytorch.org/whl/cu121强制拉取CUDA 12.1兼容的PyTorch包，若你的宿主机CUDA版本为11.8，可替换为cu118。
• 安装依赖库是为了避免容器运行时出现CUDA动态库缺失的错误。

解决方案2：Conda环境的正确安装方式 #

使用Conda时，需明确指定CUDA版本并通过官方channel安装：

# 创建Python 3.11环境
conda create -n py311-torch python=3.11 -y
conda activate py311-torch
# 安装PyTorch+CUDA 12.1（版本可按需调整）
conda install pytorch torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia -y

验证安装结果 #

安装完成后，运行以下代码确认CUDA是否正常启用：

import torch
print(torch.cuda.is_available())  # 预期输出：True
print(torch.version.cuda)         # 预期输出：对应安装的CUDA版本号（如12.1）

内容的提问来源于stack exchange，提问作者Sishaar Rao