案例分享 | 腾讯互娱基于 CPU 环境的分布式 YOLOv3 实现

文 / 腾讯 IEG 图灵实验室 张力柯 王建东

针对现实生产环境中具有大量 CPU 资源而 GPU 奇缺的现状，并出于充分利用现有 Kubernetes 的目的，我们基于 Uber 的 Horovod 实现了分布式训练框架，并且可以部署在内部 Kubernetes 平台上，通过 CPU scale 来实现机器学习模型训练，达到在 GPU 不足的情况下，通过 CPU scale 来实现模型训练，降低模型训练时间和提高算法同学模型验证效率的目标。

本文我们将主要介绍一下整体架构设计，YOLO3 的分布式算法实现过程和单机版结果对比。同时，感谢腾讯互娱 Turing Lab 全体同学在研发期间各自付出的努力。

序

最近，YOLOv3 的创始人突然宣布停止 YOLO 相关开发，但 YOLO 作为目前最经典的快速目标检测算法，仍然是各种工业应用产品中的首选。在我们的自动化检测服务中，对 UI、游戏内目标及视频内容的识别，大多是基于 YOLO 模型实现，在魂斗罗游戏中目标的识别应用如下图：

图 1  魂斗罗游戏中识别角色

YOLO 的实现，目前大多是基于原生 Darknet 的 C 语言版本 (https://pjreddie.com/darknet/)。该实现在单机环境下相当高效，无论是 training 还是 inference 都能完全满足实时监测的需求。然而其问题在于，当我们想提供目标检测的在线服务时，基于 darknet 的原生模型就失去了在线 serving 的环境，是无法直接部署为在线服务使用的。因此如果要将 YOLO 作为大规模在线服务部署，其实现方案就必须另行设计。我们希望采用通用的业界机器学习框架，同时对于产品级别的机器学习模型训练，如何实现分布式训练是必须考虑的。

架构设计

对业界分布式训练框架调研后，我们打算搭建一个基于 yolo3 分布式训练框架，主要是要结合我们的技术架构和基础设施来实现，同时，我们还希望满足以下几个需求：

• 基于业界标准开源方案如 Keras/TensorFlow
• 能够支持 CPU 训练，而不是必须 GPU 支持
• 能够支持数据并行的分布式训练

• 能够转换为 TensorFlow Serving 所支持模型

基于上面目标，在多种方案比较后，我们选择了 Kubernetes + Horovod + Keras + Tensorflow Serving 的实现，其整体流程如下：

图 2  基于 Horovod CPU 平台的分布式模型训练及部署

上图简单描述了利用 Kubernetes 和 Horovod 搭建分布式训练平台的流程，其中对最终模型的转换及利用 TensorFlow Serving 进行部署的流程在这里就不做详述。我们的重点主要是两个问题：

• 如何搭建分布式的 Horovod CPU 训练平台

• 以 YOLOv3 为例，如何修改代码，在 Keras 中实现模型训练的数据并行

YOLOv3 分布式算法实现和验证

1. 基于 Horovod CPU 环境的分布式训练平台

环境搭建：Kubernetes + Horovod

我们基于公司内部 Kubernetes 平台的扩容机制来实现分布式训练，这样我们可以充分利用 Kubernetes 对 CPU 核的灵活调度机制。

1.1 基础通用镜像

步骤：安装openmpi --> 安装python3.5 --> 安装tensorflow --> 安装 Horovod --> 安装keras框架

下面是 docker 构建镜像代码：

RUN mkdir /tmp/openmpi && \
    cd /tmp/openmpi && \
    wget https://www.open-mpi.org/software/ompi/v4.0/downloads/openmpi-4.0.0.tar.gz && \
    tar zxf openmpi-4.0.0.tar.gz && \
    cd openmpi-4.0.0 && \
    ./configure --enable-orterun-prefix-by-default && \
    make -j $(nproc) all && \
    make install && \
    ldconfig && \
    rm -rf /tmp/openmpi
 
#python3
#RUN apt-get install openssl-devel bzip2-devel expat-devel gdbm-devel readline-devel sqlite-deve
RUN wget https://www.python.org/ftp/python/3.5.1/Python-3.5.1.tgz
RUN tar -zxvf Python-3.5.1.tgz
RUN mv Python-3.5.1 /usr/local/
RUN cd /usr/local/Python-3.5.1 && \
	./configure && \
	make && \
	make install &&\
	mv /usr/bin/python /usr/bin/python.2 && \
	ln -sf /usr/local/bin/python3 /usr/bin/python
 
#python3-pip
RUN curl https://bootstrap.pypa.io/get-pip.py -o get-pip.py && \
	python get-pip.py && \
	rm -rf get-pip.py && \
	pip3  install --upgrade pip
 
RUN pip3 install tensorflow==1.4.0
 
#numpy setup
RUN pip3 uninstall numpy -y && \
	pip3 install numpy==1.16.4 &&\
  	pip3 install tensorflow
 
#horovod setup
RUN pip3 install --no-cache-dir horovod
 
#keras setup
RUN pip3 install keras==2.1.5

1.2 docker 的 ssh 交互

为了后期启动分布式 Horovod 程序，需要多个 docker 之间可以 ssh 相互登录，为了满足这个需求，我们在基础镜像中放入一个通用的 ssh 的公钥，每次启动服务自启动 ssh 服务，就可以相互登录，完成对多机的操作控制。

# Install OpenSSH for MPI to communicate between containers
RUN apt update
RUN apt-get install -y --no-install-recommends openssh-client openssh-server && \
	mkdir -p /var/run/sshd
 
# Allow OpenSSH to talk to containers without asking for confirmation
RUN	cat /etc/ssh/ssh_config | grep -v StrictHostKeyChecking > /etc/ssh/ssh_config.new && \
	echo "    StrictHostKeyChecking no" >> /etc/ssh/ssh_config.new && \
	mv /etc/ssh/ssh_config.new /etc/ssh/ssh_config

2. YOLOv3 的分布式实现

YOLOv3 基于 Keras (及部分 TensorFlow) 的实现可参考Github (https://github.com/qqwweee/keras-yolo3)，我们在这里重点对如何对其进行分布式实现进行讨论。

2.1 数据并行加载

这里用到了 Keras 的 fit_generator 函数来实现分布式加载，并且定义了自己的 data 生成器，代码如下：

class DataGenSequence(keras.utils.Sequence):
    def __init__(self, annotation_list, batch_size, input_size, step_num, anchors, num_classes):
        self.annotation_list = annotation_list
        self.item_num = len(self.annotation_list)
        self.batch_size = batch_size
        self.input_size = input_size
        self.step_num = step_num
        self.anchors = anchors
        self.num_classes = num_classes
 
    def __len__(self):
        return self.step_num
 
    def on_epoch_end(self):
        np.random.shuffle(self.annotation_list)
 
    def __getitem__(self, idx):
        img_data = []
        box_data = []
        for i in range(idx * self.batch_size, (idx + 1) * self.batch_size):
            ith = i % self.item_num
            img, box = utils.get_random_data(self.annotation_list[ith], self.input_size, random=True)
            img_data.append(img)
            box_data.append(box)
        img_data = np.array(img_data)
        box_data = np.array(box_data)
        y_true = yolomodel.preprocess_true_boxes(box_data, self.input_size, self.anchors, self.num_classes)
        return [img_data] + y_true, np.zeros(self.batch_size)

训练样本集和验证集如下：

data_gen_train = datagen_sequence.DataGenSequence(
            annotation_list=train_list,
            batch_size=self.__batch_size,
            input_size=self.__input_size,
            step_num=train_steps,
            anchors=self.__anchors,
            num_classes=self.__num_classes
        )
        data_gen_val = datagen_sequence.DataGenSequence(
            annotation_list=valid_list,
            batch_size=self.__batch_size,
            input_size=self.__input_size,
            step_num=valid_steps,
            anchors=self.__anchors,
            num_classes=self.__num_classes
        )

训练函数入口如下:

model.fit_generator(
            generator=data_gen_train,
            steps_per_epoch=train_steps,
            validation_data=data_gen_val,
            validation_steps=valid_steps,
            epochs=self.__max_epochs,
            initial_epoch=0,
            callbacks=callbacks,
            # use_multiprocessing=True,
            # workers=1
        )

2.2 分布式环境初始化

hvd.init()
hvd_size = hvd.size()
# Horovod: pin GPU to be used to process local rank (one GPU per process)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.gpu_options.visible_device_list = str(hvd.local_rank())
 K.set_session(tf.Session(config=config)

2.3 梯度计算

回调函数和优化器调整如下：

# horovod call_back
        callbacks += [
            # Horovod: broadcast initial variable states from rank 0 to all other processes.
            # This is necessary to ensure consistent initialization of all workers when
            # training is started with random weights or restored from a checkpoint.
            hvd.callbacks.BroadcastGlobalVariablesCallback(0),
            hvd.callbacks.MetricAverageCallback(),
            hvd.callbacks.LearningRateWarmupCallback(warmup_epochs=5, verbose=1),
            keras.callbacks.ReduceLROnPlateau(patience=10, verbose=1),
        ]
        
        callbacks.append(ModelModifyCallback(
           model_state_changing_epoch_list=[
               first_stage_epochs,
               second_stage_epochs,
               third_stage_epochs,
           ],
           learning_rate_list=[
               self.__learning_rate * 0.1 * hvd_size,
               self.__learning_rate * 0.01 * hvd_size,
               self.__learning_rate * 0.001 * hvd_size,
 
           ],
           trainable_changing_epoch=first_stage_epochs,
           loss=loss,
           opt = hvd.DistributedOptimizer
       ))

阶段梯度计算代码如下：

class ModelModifyCallback(keras.callbacks.Callback):
    """
    用于在训练途中修改模型状态
    """
    def __init__(self, model_state_changing_epoch_list, learning_rate_list, trainable_changing_epoch, loss, opt ):
        self.model_state_changing_epoch_list = model_state_changing_epoch_list
        self.learning_rate_list = learning_rate_list
        self.trainable_changing_epoch = trainable_changing_epoch
        self.loss = loss
        self.opt = opt
        super(ModelModifyCallback, self).__init__()
 
    def on_epoch_begin(self, epoch, logs=None):
        for i in range(len(self.model_state_changing_epoch_list)):
            if epoch == self.model_state_changing_epoch_list[i]:
                if epoch == self.trainable_changing_epoch:
                    logger.info("changing trainable weights")
                    for l in range(len(self.model.layers)):
                        self.model.layers[l].trainable = True
                logger.info("recompiling model for training stage changing (%d stage)" % (i + 1))
                
                # Horovod: adjust learning rate based on number of GPUs.
                opt = keras.optimizers.Adam(lr=self.learning_rate_list[i])
                 
                # Horovod: add Horovod Distributed Optimizer.
                opt = self.opt(opt)
 
                self.model.compile(
                    optimizer=opt,
                    loss=self.loss
                )

我们可以通过在任意一个 pod 中执行下面命令启动分布式训练，例如：

horovodrun -np  6  -H localhost:2，9.16.113.24:2，9.16.113.25:2 python train.py

该命令意味着我们将在 localhost, 9.16.113.24, 9.16.113.25 三台 pod 上各启动两个训练节点（共 6 个），执行 train.py 中的训练代码

3. 训练数据

3.1 基于不同样本数据量我们的统计数据

下图中样本量是图片数量*节点数量*epoch，可以看到在相近时间内，采用 10 个节点的分布式训练比单机训练的 Loss 减少了一倍，同时训练了近 5 倍的数据量。

图 1

图 2

3.2 基于相同样本 62*2*30 我们的统计

在采用同样训练样本数量的情况下，从下表可以看到，采用 6 个节点的分布式实现，所耗费时间只有单机时间的约 27%。

参考

下面是使用 2 个节点时的相关监控信息，可观察 CPU 及内存占用变化

下面为分别不同 epoch 和 pod 的截图：

Batch_size=62.       Epoch =30        Pod = 2

Batch_size=62.       Epoch =10        Pod = 4

Batch_size=62.       Epoch =10        Pod = 6 

Batch_size=62       Epoch =10        Pod = 10

4. 总结

从上面的实验可以看出，利用 Horovod 对 Keras 的有效支持（包括 TensorFlow 自带的 Keras）和 Kubernetes 的容器管理，我们成功地实现了以下目标：

• 基于 tensorflow.keras 的分布式训练
• 基于 Kubernetes 的 CPU 容器集群

• 达到明显可见、线性增长的分布式并发加速效果

结合图 2 所描述的整体模型训练部署流程， Horovod+Tensorflow(Keras)+Kubernetes 的方案帮助我们在 GPU 资源受限的情况下，充分利用现有 CPU 集群，顺利完成从训练到部署的完整 pipeline，满足了腾讯互娱内部多种视觉识别相关业务需求。

更多相关案例：

• 案例分享 | 24 分钟让 AI 跑起飞车类游戏
  

• 案例分享 | QQ 音乐应用 TensorFlow 构建 AI 赋能的音乐曲库

• 案例分享 | TensorFlow XLA 工作原理简介
  

  
  

加入案例分享，请点击 “阅读原文” 填写您的用例与相关信息，我们会尽快与你联系。