docker cuda pytorch

### 如何在 Docker 中设置并使用 CUDA 支持的 PyTorch 为了确保能够在 Docker 容器中成功运行带有 CUDA 加速功能的 PyTorch 应用程序，需要遵循特定的镜像构建和配置流程。 #### 创建自定义 Dockerfile 创建一个名为 `Dockerfile` 的文件，在其中指定基础镜像以及安装必要的依赖项： ```dockerfile FROM nvidia/cuda:11.7.0-cudnn8-devel-ubuntu20.04 # 更新包列表并安装基本工具 RUN apt-y --no-install-recommends \ python3-pip \ python3-dev \ git \ build-essential \ libsm6 \ libxext6 \ libxrender-dev \ vim \ ca-certificates \ curl \ gnupg \ lsb-release # 设置 Python 和 pip 版本 ENV PYTHON_VERSION=3.9 RUN curl -LO "https://github.com/pyenv/pyenv-installer/raw/master/bin/pyenv-installer" && \ chmod +x pyenv-installer && \ ./pyenv-installer && \ echo 'export PYENV_ROOT="$HOME/.pyenv"' >> ~/.bashrc && \ echo 'export PATH="$PYENV_ROOT/bin:$PATH"' >> ~/.bashrc && \ echo 'eval "$(pyenv init --path)"' >> ~/.bashrc && \ /bin/bash -c "source ~/.bashrc && pyenv install $PYTHON_VERSION && pyenv global $PYTHON_VERSION" # 配置虚拟环境 RUN python3 -m venv /opt/venv ENV PATH="/opt/venv/bin:$PATH" # 升级pip到最新版本 RUN pip install --upgrade pip setuptools wheel # 安装 PyTorch 及其依赖库 RUN pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu117 # 将工作目录设为/app 并复制当前项目代码至该位置 WORKDIR /app COPY . . CMD ["python", "./your_script.py"] ``` 此脚本基于官方 NVIDIA 提供的基础镜像来简化 GPU 开发环境搭建过程[^1]。通过这种方式可以确保所有必需组件都已就绪，并且能够充分利用硬件加速特性。 #### 构建与启动容器 完成上述操作之后，可以通过如下命令来进行镜像编译及容器实例化: ```shell $ docker build -t my_pytorch_cuda_image . $ docker run --gpus all -it --rm --name=my_container_name my_pytorch_cuda_image bash ``` 这里的关键参数是 `--gpus all` ，它允许访问主机上的全部GPU资源；如果只需要单个设备，则可替换为具体的编号如 `--gpus '"device=0"'`. #### 测试 PyTorch 是否能检测到 GPU 设备 进入正在运行的容器内部后执行下面这段简单的测试代码验证是否正确启用了CUDA支持： ```python import torch if __name__ == "__main__": print(f"CUDA Available: {torch.cuda.is_available()}") device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') tensor_example = torch.rand((2, 2), device=device) print(tensor_example) ``` 当输出显示 `"CUDA Available: True"` 表明已经成功设置了具有CUDA能力的PyTorch环境[^2].