本博客主要介绍如何在容器里面使用pytorch进行推理,训练,同时用上GPU。
1. 前置条件,安装好docker。
2. 安装nvidia-container-toolkit
参考官方文档:
Installing the NVIDIA Container Toolkit — NVIDIA Container Toolkit 1.15.0 documentation
本教程使用的是yum安装
2.1 配置仓库
curl -s -L https://nvidia.github.io/libnvidia-container/stable/rpm/nvidia-container-toolkit.repo | \sudo tee /etc/yum.repos.d/nvidia-container-toolkit.repo
2.2 启用experimental packages (可选)
sudo yum-config-manager --enable nvidia-container-toolkit-experimental
2.3 正式安装
这一步速度超级慢,但是多尝试几次还是能成功
sudo yum install -y nvidia-container-toolkit
3 docker配置
sudo nvidia-ctk runtime configure --runtime=docker
4. 重启docker
sudo systemctl restart docker
5. 构建镜像
本教程中,启动了一个fastapi服务,服务提供了一个接口使用pytroch进行模型推理
注意pytorch版本必须要和你显卡驱动版本配对
# syntax=docker/dockerfile:1
FROM pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtimeWORKDIR /appCOPY requirements.txt requirements.txtRUN pip install --no-cache-dir -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simpleENTRYPOINT ["fastapi"]CMD ["--help"]
6. 启动容器
注意--gpus all 选项,这个参数指定容器可以使用宿主机所有的GPU
docker run -d -p 9999:80 -v $(pwd):/app --gpus all --name chatcls cjq/chatcls:v1 run main.py --port 80