docker配置gpu版本的pytorch

### 如何在Docker容器中安装和配置支持GPU的PyTorch #### 安装NVIDIA驱动程序和支持工具 为了使Docker能够访问主机上的NVIDIA GPU，在宿主机上需先安装相应的NVIDIA驱动程序以及`nvidia-container-toolkit`来提供必要的运行时支持[^1]。 ```bash distribution=$(. /etc/os-release;echo $ID$VERSION_ID) curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add - curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list sudo apt-get update && sudo apt-get install -y nvidia-container-toolkit sudo systemctl restart docker ``` #### 配置Docker以使用NVIDIA Runtime 完成上述软件包的安装之后，还需确保Docker默认采用NVIDIA作为其容器运行时环境之一。这一步骤可通过命令行执行如下操作： ```bash sudo nvidia-ctk runtime configure --runtime=docker ``` 此命令会自动修改Docker守护进程配置文件(/etc/docker/daemon.json)，并添加"NVIDIA"到可用的runtimes列表里[^3]。 #### 创建带有GPU支持的PyTorch Dockerfile 创建一个新的名为`Dockerfile`的文本文件，并向其中加入以下内容用于构建一个基于官方PyTorch镜像的新镜像，该镜像已预先集成了CUDA和cuDNN的支持: ```dockerfile FROM pytorch/pytorch:latest-gpu WORKDIR /app COPY . . CMD ["python", "your_script.py"] ``` 这里选择了最新的稳定版PyTorch GPU镜像作为基础层，可以根据具体需求调整标签版本号；同时设置了工作目录并将当前项目复制进去以便后续加载Python脚本运行。 #### 构建与启动容器 保存好之前编辑好的Dockerfile后即可通过下面两条指令分别完成新自定义镜像的建立过程及其对应容器实例化的工作： ```bash docker build -t my-pytorch-app . docker run --gpus all -it --rm my-pytorch-app ``` 第一条命令负责依据指定路径下的Dockerfile生成目标应用专属的Docker镜像；第二条则用来创建临时性的交互模式容器(`-it`)，并且允许它利用所有的物理GPU资源(`--gpus all`)来进行计算加速处理[^2]。 #### 测试PyTorch是否能正常识别GPU设备 一旦进入了正在运行中的容器内部终端界面以后就可以尝试编写简单的测试代码片段确认一切设置无误： ```python import torch if __name__ == "__main__": device = 'cuda' if torch.cuda.is_available() else 'cpu' print(f'This machine uses {device} for computations.') # Create a tensor on the selected device and perform some operations. x = torch.rand(5, 5).to(device) y = (x @ x.T).sum() print(y.item()) ``` 这段小程序将会输出所选硬件平台名称（CPU或CUDA），接着还会展示一次矩阵乘法运算的结果证明整个流程已经正确搭建完毕。