[论文解读] Training Kinetics in 15 Minutes: Large-scale Distributed Training on Videos

Ji Lin, Chuang Gan+1 more|arXiv (Cornell University)|Oct 1, 2019

Human Pose and Action Recognition参考文献 23被引用 9

一句话总结

本论文提出一种硬件感知的视频模型设计方法，采用时间移位模块（TSM）实现在1,536块GPU上的可扩展、高吞吐量分布式视频训练。通过最小化浮点运算量（FLOPs）、减少输入帧数并缩小模型尺寸，TSM在1,536块GPU上相较I3D模型实现1.6倍和2.9倍的训练加速，仅用14分13秒即完成Kinetics数据集的训练，top-1准确率达到74.0%。

ABSTRACT

Deep video recognition is more computationally expensive than image recognition, especially on large-scale datasets like Kinetics [1]. Therefore, training scalability is essential to handle a large amount of videos. In this paper, we study the factors that impact the training scalability of video networks. We recognize three bottlenecks, including data loading (data movement from disk to GPU), communication (data movement over networking), and computation FLOPs. We propose three design guidelines to improve the scalability: (1) fewer FLOPs and hardware-friendly operator to increase the computation efficiency; (2) fewer input frames to reduce the data movement and increase the data loading efficiency; (3) smaller model size to reduce the networking traffic and increase the networking efficiency. With these guidelines, we designed a new operator Temporal Shift Module (TSM) that is efficient and scalable for distributed training. TSM model can achieve 1.8x higher throughput compared to previous I3D models. We scale up the training of the TSM model to 1,536 GPUs, with a mini-batch of 12,288 video clips/98,304 images, without losing the accuracy. With such hardware-aware model design, we are able to scale up the training on Summit supercomputer and reduce the training time on Kinetics dataset from 49 hours 55 minutes to 14 minutes 13 seconds, achieving a top-1 accuracy of 74.0%, which is 1.6x and 2.9x faster than previous 3D video models with higher accuracy. The code and more details can be found here: http://tsm-hanlab.mit.edu.

研究动机与目标

• 解决由于高计算、I/O和通信需求导致的视频识别训练可扩展性挑战。
• 识别分布式视频训练中的三个关键瓶颈：计算FLOPs、数据加载I/O和通信带宽。
• 提出硬件感知的模型设计准则，以提升计算、数据加载和网络效率方面的可扩展性。
• 实现在超级计算机上以分钟级而非天级完成大规模视频模型（如Kinetics）的端到端训练。
• 证明模型架构选择可显著缩短训练时间，且不损失准确性。

提出的方法

• 提出三项设计准则：通过硬件友好的算子减少FLOPs，通过最小化输入帧数减少I/O，通过减小模型尺寸降低通信开销。
• 引入时间移位模块（TSM），一种零FLOPs、零参数的算子，通过在2D卷积神经网络中跨时间维度移动通道实现时间建模。
• 采用直接的主干网络设计，使用8帧输入，相比需要大量帧的池化设计，显著降低I/O负载。
• 在Summit超级计算机上实现基于同步SGD的分布式训练，最多使用1,536块GPU，小批量大小为12,288个视频片段（98,304张图像）。
• 通过梯度累积模拟在3,072块和6,144块GPU上的训练，同时保持准确性评估。
• 通过利用高度优化的2D卷积而非效率较低的3D卷积，提升硬件利用率，优化训练效率。

实验结果

研究问题

• RQ1 计算、数据加载和通信瓶颈如何影响大规模视频训练的可扩展性？
• RQ2 模型架构设计能否显著提升视频数据集上的训练吞吐量和可扩展性？
• RQ3 在分布式视频训练中，减少FLOPs、输入帧数和模型尺寸在多大程度上能提升训练效率？
• RQ4 在1,536块GPU上，是否能实现子15分钟的Kinetics数据集训练时间，同时保持高准确率？
• RQ5 TSM模块如何在2D CNN框架中实现零FLOPs和零参数的有效时间建模？

主要发现

• 在1,536块GPU上，基于TSM的训练仅用14分13秒即完成Kinetics训练，训练时间从49小时55分钟缩短至15分钟以内。
• TSM模型在Kinetics数据集上达到74.0%的top-1准确率，尽管训练速度显著提升，但性能与或优于以往的I3D模型。
• TSM训练在1,536块GPU上达到80%的可扩展效率，表明其在硬件资源上具有出色的线性可扩展性。
• 模型在小批量大小达12,288个视频片段（98,304张图像）时仍保持准确率，性能无下降。
• 在相同硬件条件下，TSM的训练吞吐量分别比I3D 3×3×3和I3D 3×1×1模型高出1.6倍和2.9倍。
• 梯度累积模拟显示，在3,072块和6,144块GPU时出现准确率下降，表明存在稳定大规模训练的实际上限。