Docker - 制作PyTorch+YOLO+Jupyter模型训练环境镜像教程（深度学习环境）

作者：hangge | 2024-09-26 08:39

要训练 AI 模型，我们除了直接可以在物理机上训练数据集外，也可在 Docker 容器中进行训练。并且由于 Docker 的可移植性，能让我们十分方便地进行安装部署、环境配置、备份还原等操作。本文分别通过 docker commit 和 Dockerfile 这两种方式演示如何制作一个 PyTorch + YOLOv5 + Jupyter 的算法镜像，以及如何使用这个镜像。

方式一：使用 docker commit 命令构建镜像

1，下载基础镜像

（1）首先我们要拉取一个 PyTorch 的镜像作为基础镜像，访问 docker hub 页面（点击访问），搜索 pytorch

（2）根据自己的情况选择需要的版本，比如我这里选择 1.12.1-cuda11.3-cudnn8-devel

devel 与 runtime 两个版本的区别：

devel 是说涵盖了开发所需的所有工具，包含编译、debug 等，以及编译需要的头文件、静态库。
runtime 是说只涵盖了运行环境的最小集合，例如动态库等。一般来说，使用该版本也是没问题的。

（3）复制该版本右侧的命令，然后在终端中执行即可将镜像拉取到本地。

docker pull pytorch/pytorch:1.12.1-cuda11.3-cudnn8-devel

2，启动容器

（1）首先执行如下命令创建一个容器：

docker run -it -d --name pytorch pytorch/pytorch:1.12.1-cuda11.3-cudnn8-devel

（2）接着执行如下命令进入该容器：

docker attach pytorch

3，安装常用软件包

（1）在容器内执行如下命令建立软件包缓存：

apt update

（2）然后执行如下命令安装常用软件以及必要的依赖库：

apt install -y vim git libgl1-mesa-glx libglib2.0-dev

4，安装 yolov5

（1）执行如下命令将 yolov5 项目下载到容器中：

注意：如果 yolov5 项目无法从 github 上 clone 下来的话，可以考虑改用国内的镜像仓库（点击查看）

git clone https://github.com/ultralytics/yolov5

（2）接着执行如下命令安装 yolov5 相关依赖：

提示：为了加快 pip install 安装速度，这里我临时指定了国内的镜像源。

pip install -r ./yolov5/requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

5，安装 Jupyter

（1）首先执行如下命令安装 Jupyter：

pip install jupyter -i https://pypi.tuna.tsinghua.edu.cn/simple

（2）接着执行如下命令生成配置文件：

jupyter notebook --generate-config

（3）然后编辑刚才创建的配置文件：

vi /root/.jupyter/jupyter_notebook_config.py

（4）在配置文件中添加如下内容后保存退出：

c.NotebookApp.ip = "0.0.0.0"
c.NotebookApp.open_browser = False
c.NotebookApp.port = 8888
c.NotebookApp.token = ""
c.NotebookApp.allow_root = True
c.NotebookApp.notebook_dir = r"/"

6，构建镜像

执行 exit 命令退出容器后，执行如下命令即可将该容器生成镜像：

注意：直接 commit 通常来说是没有 entrypoint 的，但我们需要在容器启动时能够自动运行 Jupyter，那么可以借助 --change 参数来实现。

docker commit --change='ENTRYPOINT ["jupyter", "notebook", "--allow-root"]' pytorch pytorch-yolov5-jupyter:v1.0

方式二：使用 Dockerfile 构建镜像

1，创建一个 Dockerfile 文件

（1）Dockerfile 其实是一个文本文件，记录了镜像构建的所有步骤。我们可以通过 vi 命令创建它。

vi Dockerfile

（2）Dockerfile 里的内容如下：

注意：

为了加快 pip install 安装速度，这里我临时指定了国内的镜像源。
如果 yolov5 项目无法从 github 上 clone 下来的话，可以考虑改用国内的镜像仓库（点击查看）

FROM pytorch/pytorch:1.12.1-cuda11.3-cudnn8-devel
# 设置作者信息
MAINTAINER hangge.com
# 建立软件包缓存(|| true 作用是避免部分缓存更新失败时不会中断整个镜像构建)
RUN apt update -y || true
# 安装常用软件以及必要的依赖库
RUN apt install -y vim git libgl1-mesa-glx libglib2.0-dev
# 设置工作目录
WORKDIR /workspace
# 将 yolov5 项目下载到镜像中
Run git clone https://github.com/ultralytics/yolov5
# 安装 yolov5 相关依赖
RUN pip install -r ./yolov5/requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple --trusted-host pypi.tuna.tsinghua.edu.cn
# 安装 jupyter
RUN pip install jupyter -i https://pypi.tuna.tsinghua.edu.cn/simple --trusted-host pypi.tuna.tsinghua.edu.cn
# 生成 jupyter 配置文件
RUN jupyter notebook --generate-config
# 修改 jupyter 配置文件
RUN sed -i '1i\c.NotebookApp.ip = "0.0.0.0"\nc.NotebookApp.open_browser = False\nc.NotebookApp.port = 8888\nc.NotebookApp.token = ""\nc.NotebookApp.allow_root = True\nc.NotebookApp.notebook_dir = r"/"\n' /root/.jupyter/jupyter_notebook_config.py
# 声明需要暴露的端口
EXPOSE 8888
# 配置容器启动后执行的命令
#ENTRYPOINT ["./entrypoint.sh"]
ENTRYPOINT ["jupyter", "notebook", "--allow-root"]

2，构建镜像

（1）Dockerfile 文件创建好后，运行 docker build 命令即可开始构建镜像：

参数说明：

-t 表示将新镜像命名为 hangge_server
末尾的 . 指明 build context 为当前目录
Docker 默认会从 build context 中查找 Dockerfile 文件，我们也可以通过 -f 参数指定 Dockerfile 的位置。

docker build -t pytorch-yolov5-jupyter:v1.0 .

（2）上面执行后会显示出详细的构建过程。

附：镜像使用

1，运行容器

（1）镜像构建完毕后，我们执行如下命令进行运行一个容器：

注意：启动命令增加 --gpus device=all 参数表示使用 GPU。当然服务器也需要安装 docker 使用 GPU 的 nvidia 的工具包，具体可以参考我之前写的文章（点击查看）

docker run -it -d -p 8888:8888 --gpus device=all --name my-pytorch pytorch-yolov5-jupyter:v1.0

（2）执行 docker ps 命令可以看到容器已经运行：

2，训练测试

（1）我们可以通过命令进入容器来进行相关训练操作，但由于我们镜像还安装了 Jupyter，所以我们也可以通过浏览器访问：

（2）点击右侧菜单的 New -> Terminal 创建一个终端命令窗口：

（3）在终端中进入 yolov5 文件夹，然后执行 python detect.py 命令，代码会自动下载模型并对项目中自带的两张图片进行目标检测：

（4）打开终端中提示的结果保存目录，可以看到确实有两张照片：

（5）打开照片可以看到目标也被准确标注出来，说明环境是正常可用的。

Docker / K8s