Docker部署AI模型GPU加速实战教程

为什么要在Docker中用GPU跑AI模型？

很多新手在部署AI模型时会被环境配置卡住——Python版本、CUDA、cuDNN、驱动版本，任意一个不匹配就会报错。Docker + GPU加速的组合能让你把整个运行环境打包成一个镜像，无论换到哪台机器，只要装上Docker和NVIDIA驱动，一条命令就能启动模型推理。

下面我会从零开始，带你把一个AI模型跑在Docker容器里，并确认GPU已经被正确调用。

准备工作：检查你的硬件和驱动

在动手之前，先确认你的服务器满足三个条件：

• 有一块NVIDIA显卡（型号不限，支持CUDA即可）


• 已安装Linux系统（Ubuntu 20.04/22.04最省心）


• 确保你能用sudo执行命令

第一步：安装NVIDIA驱动

如果你还没有驱动，运行以下命令自动安装推荐版本：

sudo ubuntu-drivers autoinstall

sudo reboot

重启后，用nvidia-smi查看驱动信息，出现当前驱动版本和显卡列表就说明驱动正常。

如果命令找不到，先手动安装：sudo apt install nvidia-driver-535（535是常见稳定版，具体版本根据你的显卡选择）。

第二步：安装Docker和NVIDIA Container Toolkit

Docker官方安装脚本：

curl -fsSL https://get.docker.com | sudo sh

sudo usermod -aG docker $USER

newgrp docker  # 让当前用户能免sudo执行docker

然后安装NVIDIA容器工具包，让Docker能读取到宿主的GPU：

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 update && sudo apt install -y nvidia-container-toolkit

sudo systemctl restart docker

完成后，

执行docker run --rm --gpus all nvidia/cuda:

11.8-base nvidia-smi，

如果能成功输出nvidia-smi信息，

就说明Docker已经能调用GPU了。

部署AI模型：拉取镜像并启动容器

这里我用一个经典的PyTorch推理模型举例，你可以替换成你自己的模型文件。

1. 选择基础镜像

推荐使用PyTorch官方提供的GPU镜像，比如：

docker pull pytorch/pytorch:1.13.1-cuda11.6-cudnn8-runtime

这个镜像已经包含了CUDA 11.6和cuDNN 8，适合大多数PyTorch模型。

2. 编写Dockerfile（如果有定制需求）

如果模型依赖一些额外Python包，创建一个Dockerfile：

FROM pytorch/pytorch:1.13.1-cuda11.6-cudnn8-runtime

WORKDIR /app

COPY requirements.txt .

RUN pip install --no-cache-dir -r requirements.txt

COPY model.pth inference.py ./   # 把你的模型文件和推理脚本放进来

CMD ["python", "inference.py"]

简单场景可以直接用现有镜像，不需要自己写Dockerfile。

3. 运行容器并挂载GPU

假设你的模型文件和推理脚本都在/home/user/my_model目录，执行：

docker run --rm --gpus all -v /home/user/my_model:/workspace pytorch/pytorch:1.13.1-cuda11.6-cudnn8-runtime python /workspace/inference.py

关键参数说明：

• --gpus all：让容器使用所有GPU，你也可以指定某个GPU（如--gpus 'device=0'）


• -v：将宿主目录挂载到容器内，方便修改代码和模型

验证GPU是否真正生效

模型跑完后，你肯定想知道GPU到底有没有被用上。

两种方法：

方法1：在容器内打印GPU信息

在你的inference.py中加入以下代码：

import torch

print("CUDA available:", torch.cuda.is_available())

print("GPU count:", torch.cuda.device_count())

print("Current device:", torch.cuda.get_device_name(0))

运行后如果输出CUDA available: True，就说明成功。

方法2：在宿主管看GPU占用

容器运行期间，打开另一个终端，执行watch -n 1 nvidia-smi。你会看到进程列表中出现了python进程，且显存和GPU利用率在跳动。如果显示的是No running processes found，那证明模型没用到GPU，需要检查代码中是否把模型和数据都放到了.cuda()上。

高频问题与避坑指南

Q1：--gpus参数报错“Unknown runtime specified nvidia”

→ 检查nvidia-container-toolkit是否安装并重启Docker：sudo systemctl restart docker。如果还不行，检查/etc/docker/daemon.json是否包含"runtimes": {"nvidia": {...}}，没有则手动添加后重启。

Q2：容器启动后nvidia-smi显示“Failed to initialize NVML”

→ 常见原因是宿主驱动版本太旧或容器内CUDA版本不匹配。用nvidia-smi查看宿主驱动对应的最高CUDA版本，然后选择一个低于该版本的镜像。例如驱动535支持CUDA 12.2，你可以用pytorch/pytorch:2.0.1-cuda11.8-cudnn8（CUDA 11.8更低）。

Q3：镜像拉取速度极慢

→ 配置国内镜像加速器。编辑/etc/docker/daemon.json，加入：

{

"registry-mirrors": ["https://docker.mirrors.ustc.edu.cn"]

}

然后sudo systemctl restart docker再次拉取。

Q4：模型推理速度反而比CPU还慢

→ 检查代码中是否真的把张量和模型移到了GPU。常见错误：model.to('cuda')却忘记把输入数据也.cuda()。另外，如果模型很小，GPU初始化开销可能超过计算收益，这种情况属于正常现象。

总结

通过以上步骤，你已经能用一个干净的Docker容器跑AI模型，并且确认了GPU加速正常开启。

以后再迁移模型到其他服务器，只需要重复“装驱动+装Docker+拉镜像”三个动作，省去了大量环境折腾的时间。

如果你在操作中遇到本文没覆盖的报错，建议先检查nvidia-docker的官方文档，或者把错误信息贴到评论区一起讨论。