pytorch在docker里面使用GPU

本博客主要介绍如何在容器里面使用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-runtime

WORKDIR /app

COPY requirements.txt requirements.txt

RUN pip install --no-cache-dir -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

ENTRYPOINT ["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