1. 深度视频模型的分布式训练

在视频上训练深度神经网络非常耗时。例如,使用一台配有 8 块 V100 GPU 的服务器,在 Kinetics400 数据集(包含 24 万个 10 秒短视频)上训练最先进的 SlowFast 网络 [Feichtenhofer18] 需要超过 10 天。缓慢的训练导致研究周期漫长,对新手和学生研究视频相关问题不友好。

使用分布式训练是一个自然的选择。将庞大的计算量分散到多台机器上可以极大地加速训练。然而,关于视频理解的开源 Github 仓库中只有少数支持分布式训练,并且通常缺乏对这一功能的文档。此外,在线上关于如何对深度视频模型进行分布式训练的信息/教程也不多。

因此,我们在此提供一个简单的教程,演示如何使用我们的代码在 Kinetics400 数据集上执行 SlowFast 模型的分布式训练。

分布式训练

我们可以通过两种方式在多个设备之间分配神经网络训练的工作负载:数据并行和模型并行。数据并行是指每个设备存储模型的完整副本。每个设备处理数据集的不同部分,并且这些设备共同更新一个共享模型。当模型非常大以至于无法容纳在设备内存中时,模型并行就很有用。在这种情况下,不同的设备被分配学习模型的不同部分的任务。在本教程中,我们将介绍如何以数据并行的方式在跨机器分布的设备上训练模型。

在分布式训练中,有一些关键概念,例如 server(服务器)、worker(工作节点)、scheduler(调度器)、kvstore 等。Server 是一个存储模型参数或梯度的节点,并与 workers 通信。Worker 是实际在训练样本批次上执行训练的节点。在处理每个批次之前,workers 从 servers 拉取权重。Workers 在每个批次之后也会将梯度发送到 servers。Scheduler 用于设置集群以便每个节点通信,并且整个集群中只有一个 scheduler。Kvstore,即键值存储,是用于多设备训练的关键组件。它存储模型参数和优化器,用于接收梯度并更新模型。为了保持本教程的简洁性,我不会深入细节。读者可以参考 MXNet 官方文档以获取更多信息。

如何使用我们的代码以分布式方式训练 SlowFast 模型?

为了执行分布式训练,您需要 (1) 准备集群;(2) 安装 MXNet;以及 (3) 准备您的代码和数据。

我们需要一个集群,其中每个节点都可以相互通信。第一步是为每台机器生成 ssh 密钥。为了更好地说明,假设我们有四台机器:node1、node2、node3 和 node4。我们将每台机器都用作 server 和 worker。请注意,我们也可以使用专用的 CPU 机器作为 servers。

首先,ssh 到 node1 并输入

ssh-keygen -t rsa

按照默认设置,您将拥有一个名为 id_rsa 的文件和一个名为 id_rsa.pub 的文件,两者都在 ~/.ssh/ 文件夹下。id_rsa 是私有 RSA 密钥,而 id_rsa.pub 是其公钥。

其次,将 node1 的这两个文件(id_rsaid_rsa.pub)复制到所有其他机器。对于每台机器,您也会在 ~/.ssh/ 文件夹下找到一个 authorized_keys 文件。将 authorized_keys 的内容追加到 id_rsa.pub 中。这一步将确保集群中的所有机器都能相互通信。

在继续下一步之前,最好进行一些健全性检查,以确保通信良好。例如,如果您可以成功 ssh 到其他机器,则表示它们现在可以相互通信了。您就可以继续了。如果在 ssh 过程中出现任何错误,您可以使用 -vvv 选项来获取详细信息以进行调试。

一旦您准备好集群,就可以安装 MXNet 了。MXNet 提供了一个脚本 tools/launch.py,可以轻松地在具有 ssh、mpi、sge 或 yarn 的集群上启动分布式训练。让我们首先在所有机器上安装 MXNet。

pip install mxnet-cu100

有关更多安装选项(即不同版本的 MXNet 或 CUDA),请查阅 GluonCV 安装指南以获取更多信息。另外提一点,clush 是一个很好的集群设置工具。它可以用于在主机集群上并发执行命令,但我们在此不作详细介绍。

接下来,我们需要启动作业的脚本,让我们也克隆一下仓库。

git clone https://github.com/apache/incubator-mxnet.git --recursive

我们需要的脚本位于文件夹 tools 下,名为 launch.py。请注意,这个脚本可以放在任何节点上,因为我们只需要它来启动作业。

现在是时候在每个节点上准备您的代码和数据了。代码需要位于每台机器上的同一个目录中,以便单个命令可以在多台机器上工作。让我们克隆 GluonCV 仓库并安装它,

git clone https://github.com/dmlc/gluon-cv.git
cd gluon-cv
pip install -e .

类似地,数据也需要位于每台机器上的同一个路径中,以便数据加载器知道在哪里找到数据。

最后,我们可以开始训练了。让我们输入下面的命令来启动一个有 4 台机器参与的训练作业。

../incubator-mxnet/tools/launch.py -n 4 -H host.txt --launcher ssh \
python ./scripts/action-recognition/train_recognizer.py --dataset kinetics400 \
--data-dir ~/.mxnet/kinetics400/train --val-data-dir ~/.mxnet/kinetics400/val \
--train-list ~/.mxnet/kinetics400/train.txt --val-list ~/.mxnet/kinetics400/val.txt \
--dtype float32 --mode hybrid --prefetch-ratio 1.0 --kvstore dist_sync_device \
--model slowfast_4x16_resnet50_kinetics400 --slowfast --slow-temporal-stride 16 \
--fast-temporal-stride 2 --video-loader --use-decord --num-classes 400 --batch-size 8 \
--num-gpus 8 --num-data-workers 32 --input-size 224 --new-height 256 --new-width 340 \
--new-length 64 --new-step 1 --lr-mode cosine --lr 0.4 --momentum 0.9 --wd 0.0001 \
--num-epochs 196 --warmup-epochs 34 --warmup-lr 0.01 --scale-ratios 1.0,0.8 \
--save-frequency 10 --log-interval 50 --logging-file slowfast_4x16.log --save-dir ./checkpoints

这里,host.txt 文件包含所有机器的私有 IP 地址,例如,

123.123.1.123
123.123.2.123
123.123.3.123
123.123.4.123

--kvstore dist_sync_device 是指每个节点使用多个 GPU 时,该模式在 GPU 上聚合梯度并更新权重。还有其他模式,例如 dist_sync、dist_async 等。要了解更多详情,请查阅本教程

另一个需要注意的地方是学习率。对于单节点训练,我们将学习率设置为 0.1。然而,对于多节点训练,由于我们有四台机器,我们将其增加到 0.4。这是 [Goyal17] 中介绍的一种良好实践,称为线性缩放规则。当 minibatch 大小乘以 k 时,学习率也乘以 k。所有其他超参数(权重衰减等)保持不变。线性缩放规则不仅可以帮助我们使使用小 minibatch 和大 minibatch 时的准确率匹配,同样重要的是,它还能大致匹配它们的训练曲线,从而在收敛之前实现快速调试和实验比较。

如果一切设置顺利,您将看到模型正在训练。所有打印的信息都将被捕获并发送到运行 launch.py 的机器。检查点将保存在每台机器的本地。

速度

通常,使用更多机器训练会更快,但由于通信开销,速度提升不是线性的。实际的速度提升取决于网络带宽、服务器 CPU 能力、数据加载效率等。例如,如果您在 AWS 上同一放置组中的四台 P3.16xlarge 机器上使用我们的代码,您将获得 3 倍的速度提升。使用 8 台或更多机器时,可以观察到类似的速度提升比例 (0.75)。在我们的案例中,我们使用八台 P3.16xlarge 机器训练一个 slowfast_4x16_resnet50_kinetics400 模型。训练可以在 1.5 天内完成。

常见问题

问题 1:我看到下面的错误消息,我该怎么办?

Permissions for 'id_rsa' are too open.
It is recommended that your private key files are NOT accessible by others.
This private key will be ignored.

答案:复制后,您需要通过 chmod 400 id_rsa 确保 id_rsa 的权限是正确的。

问题 2:我找不到文件 .ssh/authorized_keys

答案:只需创建一个,使用

touch ~/.ssh/authorized_keys
chmod 600 ~/.ssh/authorized_keys

参考文献

Goyal17

Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, Kaiming He. “准确、大批量 SGD:在 1 小时内训练 ImageNet。” arXiv preprint arXiv:1706.02677 (2017)。

Feichtenhofer18

Christoph Feichtenhofer, Haoqi Fan, Jitendra Malik, Kaiming He. “用于视频识别的 SlowFast 网络。” arXiv preprint arXiv:1812.03982 (2018)。

脚本总运行时间: ( 0 分钟 0.000 秒)

由 Sphinx-Gallery 生成的图库