pyorch代码封装成docker进行模型训练及通过tensorboard查看远程服务器内容器中的训练过程

1 镜像封装

yolov5的训练代码中给出了生成docker镜像的Dockerfile，如下所示：

# Start FROM Nvidia PyTorch image https://ngc.nvidia.com/catalog/containers/nvidia:pytorch
FROM nvcr.io/nvidia/pytorch:20.10-py3

# Install dependencies
RUN pip install --upgrade pip
# COPY requirements.txt .
# RUN pip install -r requirements.txt
RUN pip install gsutil

# Create working directory
# RUN mkdir -p /usr/src/app
# WORKDIR /usr/src/app
RUN mkdir -p /home/ccf/yolov5_train
WORKDIR /home/ccf/yolov5_train

# Copy contents
# COPY . /usr/src/app
COPY . /home/ccf/yolov5_train

上面的代码中，最重要的就是第一句，拉取nvidia的pyTorch镜像。这是什么？答案在这里：https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/index.html。个人的理解就是nvidia把基于pyorch进行模型训练和推理的环境全部打了一个docker镜像，这样用户只需要将自己的代码放进去组一个新的docker镜像就可以在一台新的机器上进行模型的训练和推理。当然如果需要使用GPU进行训练的话，需要宿主机上配备Nvidia显卡、安装了适配版本的显卡驱动且安装了nvidia-docker。

nvidia的pyTorch镜像内部默认包含了下面的东西，可以满足常见的模型训练和部署：
显卡驱动：

不同版本的nvidia的pyTorch镜像支持的环境版本也不同，参见下表：

需要注意的一点是，每一个版本的nvidia的pyTorch镜像都有对应的显卡驱动版本要求，如Release 20.09 is based on NVIDIA CUDA 11.0.3, which requires NVIDIA Driver release 450 or later. However, if you are running on Tesla (for example, T4 or any other Tesla board), you may use NVIDIA driver release 418.xx or 440.30. The CUDA driver’s compatibility package only supports particular drivers. For a complete list of supported drivers, see the CUDA Application Compatibility topic. For more information, see CUDA Compatibility and Upgrades.

如果需要依据环境需要使用20.09版本，但显卡驱动又不兼容的话，那么就按需进行显卡驱动升级。

不同版本的nvidia-docker使用命令也有所区别：

2 镜像的使用

以在一台桌面版ubuntu机器(一号机)上将yolov5的训练代码部署到一台ubuntu1604 server的服务器(二号机)上为例。一号机只有单显卡，可以调试代码，但不适合进行大规模训练。二号机有八块显卡，适合进行模型训练，但因为在内网环境中，直接配置pyorch的模型训练环境比较麻烦，所以需要在一号机器上将训练代码封装成docker镜像在二号机器上训练。

docker相关资料

2.1 封装docker镜像

在一号机上将训练代码中的data/coco.yaml中的训练和验证数据加载地址修改为如下所示：

# download command/URL (optional)
download: bash data/scripts/get_coco.sh

# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
train: ./datasets/train2017.txt  # 118287 images
val: ./datasets/val2017.txt  # 5000 images
#test: ../coco/test-dev2017.txt  # 20288 of 40670 images, submit to https://competitions.codalab.org/competitions/20794

# number of classes
nc: 80

# class names
names: ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
        'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
        'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
        'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
        'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
        'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
        'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
        'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
        'hair drier', 'toothbrush']

# Print classes
# with open('data/coco.yaml') as f:
#   d = yaml.load(f, Loader=yaml.FullLoader)  # dict
#   for i, x in enumerate(d['names']):
#     print(i, x)

在训练代码所在的文件夹下，执行：

#构建镜像
sudo nvidia-docker build -t yolov5_train:V1.1 .

#保存镜像为本地文件
sudo docker save -o yolov5_train:V1.1 yolov5_train_V1.1.tar

#授权
sudo chmod -R 777 yolov5_train_V1.1.tar

将镜像文件及训练数据集放到服务器上：

scp -r yolov5_train_V1.1.tar 服务器登陆用户名@服务器IP:文件在服务器上的存放地址

2.2 在二号机上使用docker镜像及通过tensorboard查看远程服务器内容器中的训练过程

因为二号机是ubuntu server的服务器，放置在机房内，需要在一号机上通过ssh远程连接。二号机就是训练代码docker镜像的宿主机。

2.2.1 远程连接

首先通过下面的方式在一号机上远程连接二号机：

ssh -L 16006:127.0.0.1:6006 username@IP

• 16006:127.0.0.1 一号机的16006端口
• 6006 username@IP: IP 二号机的IP 地址，username 二号机的登陆用户名，6006 二号机的6006端口

完整的作用是以username的身份远程登陆地址为IP的二号机，且建立ssh隧道，将二号机的6006端口转发到一号机的16006端口。在一号机和二号机之间构建ssh隧道的作用是为了通过tensorboard实时查看模型训练结果，2.2.4节会给出更加详细的解释。

因为已经执行了进行了远程连接，所以下面的操作都是在二号机器上进行操作的。

2.2.2 创建容器

#加载docker镜像
sudo nvidia-docker load -i yolov5_train_v1.1.tar

#以交互式方式创建容器
sudo nvidia-docker run -it --shm-size 8G -v 训练图像在宿主机上的位置:容器的workdir/datasets -v 容器的workdir/runs:宿主机上模型的存放位置 -p 6006:6006 yolov5_train:V1.1

• -v 指定将宿主机上的训练图像集映射到容器的workdir/datasets，datasets的名字是因为data/coco.yaml中指定的train/val数据都放在其下；

• -v 指定将训练的模型映射到宿主机上；

• -p 指定将容器的6006端口映射到宿主机的6006端口，目的也是为了可以远程通过tensorboard查看训练过程；

• –shm-size 8G 是为了解决下述bug：

  下述关于共享内存的问题参考自：https://blog.csdn.net/u013823233/article/details/101209718

  ERROR: Unexpected bus error encountered in worker. This might be caused by insufficient shared memory (shm).
  

  原因如下所示，多进程训练时需要设置足够大的共享内存存放训练数据：

  Please note that PyTorch uses shared memory to share data between processes, so if torch multiprocessing is used (e.g. for multithreaded data loaders) the default shared memory segment size that container runs with is not enough, and you should increase shared memory size either with --ipc=host or --shm-size command line options to nvidia-docker run.

  解决方案：一是创建容器时使用–shm-size指定足够大的贡献内存；二是设置torch.utils.data.DataLoader的num_workers=0，即用单个线程读取数据。

2.2.3 启动训练

进入容器内，假设采用DDP模式进行训练，执行：

python -m torch.distributed.launch --nproc_per_node 8 train.py --batch-size 256 --data data/coco.yaml --cfg models/yolov5s.yaml

至此，使用docker镜像封装的训练代码开始了模型训练，方便快捷。

2.2.4 通过tensorboard查看远程服务器内容器中的训练过程

训练是在二号机的容器内，容器内可以通过tensorboard实时反馈训练过程，但是没办法可视化的。二号机器上没有tensorboard的环境，也没有可视化的桌面。只能在一号机器上远程可视化tensorboard的训练过程。

容器的6006端口映射到宿主机(二号机)的6006端口，二号机的6006端口又和一号机的16006端口构建了ssh隧道，因此容器的6006端口呈现的内容都可以在一号机器的16006端口上可视化出来。这也就是为什么创建容器时添加-p 6006:6006，远程连接时使用ssh -L 16006:127.0.0.1:6006 username@IP的原因了。windows/Linux构建ssh隧道请参考：https://blog.csdn.net/aiynmimi/article/details/88531695，一号机器可以有多台，都可以通过构建自身某个端口和二号机器6006端口之间的ssh隧道，做到一份训练代码在跑，多个人查看训练中间过程。

剩下的工作就是在容器内启动tensorboard了：

#在一号机上新建ssh连接
ssh -L 16006:127.0.0.1:6006 username@IP

#查看已启动的容器
sudo nvidia-docker ps -a

找到正在进行模型训练的容器的ID，如xxxxxxxxxxx，执行：

#以交互式方式进入容器
sudo nvidia-docker exec -it xxxxxxxxxxx /bin/bash

#启动tensorboard
tensorboard --logdir=日志文件所在的路径

在一号机器的浏览器中输入127.0.0.1:16006，可以通过tensorboard可视化模型训练的中间过程值。

模型已经映射到了宿主机上，训练完成后可以通过scp将模型从二号机器上拷贝出来。

scp -r username@IP：模型所在路径 模型的本地存储路径

训练过程可视化和模型获取完成。