MMAction2 视频理解框架训练与测试全指南

MMAction2 视频理解框架训练与测试全指南

mmaction2 OpenMMLab's Next Generation Video Understanding Toolbox and Benchmark 项目地址: https://gitcode.com/gh_mirrors/mm/mmaction2

前言

MMAction2 是一个基于 PyTorch 的开源视频理解工具库，支持多种视频理解任务，如动作识别、时空动作检测等。本文将详细介绍如何在 MMAction2 框架下进行模型的训练和测试，涵盖从单机单卡到分布式集群的各种场景。

训练篇

单 GPU 训练

对于大多数研究者和开发者来说，单 GPU 训练是最常见的场景。MMAction2 提供了简单易用的训练脚本：

python tools/train.py configs/recognition/tsn/tsn_r50_1x1x3_100e_kinetics400_rgb.py

关键参数解析：

• --work-dir: 指定工作目录，用于保存训练日志和模型权重
• --resume: 从检查点恢复训练，支持自动恢复最新权重
• --amp: 启用自动混合精度训练，可显著减少显存占用
• --cfg-options: 动态覆盖配置文件参数，格式为 key=value

CPU 训练技巧： 虽然不推荐，但在没有 GPU 的环境下，可以通过设置环境变量强制使用 CPU：

CUDA_VISIBLE_DEVICES=-1 python tools/train.py ...

多 GPU 分布式训练

当需要加速训练或处理大规模数据集时，多 GPU 训练是必然选择。MMAction2 使用 PyTorch 的分布式数据并行(DDP)实现：

bash tools/dist_train.sh configs/...py 8 --work-dir work_dirs/exp1

环境变量控制：

• PORT: 指定通信端口，默认为 29500
• CUDA_VISIBLE_DEVICES: 控制可见 GPU 设备

多任务并行技巧： 在同一台服务器上启动多个训练任务时，需要确保：

1. 使用不同的 GPU 组
2. 设置不同的通信端口

# 任务1
CUDA_VISIBLE_DEVICES=0,1,2,3 PORT=29500 bash tools/dist_train.sh ... 4
# 任务2 
CUDA_VISIBLE_DEVICES=4,5,6,7 PORT=29501 bash tools/dist_train.sh ... 4

多机分布式训练

基础网络环境

在多台机器组成的集群中进行训练时，需要特别注意网络配置：

# 主节点(rank 0)
NNODES=2 NODE_RANK=0 PORT=29500 MASTER_ADDR=192.168.1.1 bash tools/dist_train.sh ... 8

# 从节点(rank 1) 
NNODES=2 NODE_RANK=1 PORT=29500 MASTER_ADDR=192.168.1.1 bash tools/dist_train.sh ... 8

关键环境变量：

• NNODES: 总节点数
• NODE_RANK: 当前节点序号(0-based)
• MASTER_ADDR: 主节点IP地址
• PORT: 所有节点必须使用相同端口

Slurm 集群管理

在 HPC 环境中，通常使用 Slurm 作业调度系统：

GPUS=16 GPUS_PER_NODE=8 CPUS_PER_TASK=10 \
bash tools/slurm_train.sh partition_name job_name config.py

Slurm 参数优化：

• GPUS_PER_NODE: 根据节点实际 GPU 数量设置
• CPUS_PER_TASK: 建议设置为 GPU 数量的 1-2 倍
• SRUN_ARGS: 可传递额外的 Slurm 参数

测试篇

单 GPU 测试

模型评估与训练使用相似的接口：

python tools/test.py config.py checkpoint.pth --eval top_k_accuracy

可视化功能：

• --show: 实时显示预测结果
• --show-dir: 保存可视化结果到指定目录
• --interval: 控制采样间隔
• --wait-time: 控制每帧显示时间

多 GPU 分布式测试

与训练类似，使用分布式测试可以加速大规模数据集的评估：

bash tools/dist_test.sh config.py checkpoint.pth 8 --metrics top_k_accuracy

注意事项：

1. 确保所有 GPU 使用相同版本的模型权重
2. 分布式测试时部分可视化功能可能受限

多机测试方案

基础网络测试

多机测试的配置与训练完全一致，只需替换脚本为 dist_test.sh：

# 节点0
NNODES=2 NODE_RANK=0 MASTER_ADDR=... bash tools/dist_test.sh ...

# 节点1
NNODES=2 NODE_RANK=1 MASTER_ADDR=... bash tools/dist_test.sh ...

Slurm 测试任务

对于 Slurm 集群，使用专用测试脚本：

bash tools/slurm_test.sh partition job_name config.py checkpoint.pth

性能优化建议：

1. 适当增加 CPUS_PER_TASK 加速数据加载
2. 对于大模型，增加 GPUS_PER_NODE 减少通信开销

最佳实践

1. 学习率调整：使用 --auto-scale-lr 根据实际 batch size 自动调整学习率
2. 随机种子：通过 --seed 和 --diff-rank-seed 控制实验可复现性
3. 配置覆盖：灵活使用 --cfg-options 快速实验不同超参数
4. 混合精度：训练时添加 --amp 可减少显存占用，加快训练速度

通过本文介绍的各种训练和测试方法，用户可以根据自身硬件条件和任务需求，灵活选择最适合的方案来使用 MMAction2 进行视频理解任务的研究和开发。

mmaction2 OpenMMLab's Next Generation Video Understanding Toolbox and Benchmark 项目地址: https://gitcode.com/gh_mirrors/mm/mmaction2