用TSM训练自定义动作识别模型:从UCF101格式准备到避坑调参全流程(PyTorch 1.10)
基于TSM的视频动作识别实战从数据准备到模型调优全流程解析视频动作识别作为计算机视觉领域的重要研究方向在智能监控、人机交互、运动分析等场景中具有广泛应用价值。Temporal Shift ModuleTSM通过创新的时序特征建模方式在保持2D CNN计算效率的同时实现了接近3D CNN的性能表现。本文将系统性地介绍如何利用PyTorch框架从原始视频数据开始构建完整的TSM训练流程。1. 环境配置与数据准备1.1 基础环境搭建TSM模型训练需要以下核心组件支持# 基础环境安装命令 conda create -n tsm python3.8 conda install pytorch1.10.0 torchvision0.11.0 cudatoolkit11.3 -c pytorch pip install opencv-python ffmpeg-python硬件配置建议GPU至少8GB显存如RTX 2070及以上内存16GB以上存储SSD硬盘视频数据处理对I/O性能要求较高1.2 数据集格式规范TSM默认采用UCF101数据集格式目录结构示例如下EventDataset/ ├── videos/ │ ├── Action1/ │ │ ├── v_Action1_g01_c01.avi │ │ └── v_Action1_g01_c02.avi │ └── Action2/ │ ├── v_Action2_g01_c01.avi │ └── v_Action2_g01_c02.avi └── splits/ ├── trainlist01.txt └── testlist01.txt关键命名规则gXX表示视频组别编号cXX表示同一组别下的片段编号动作类别通过子目录名称体现2. 数据预处理全流程2.1 视频帧提取技术细节使用TSM源码中的vid2img_ucf101.py脚本时需特别注意以下参数配置# 关键参数调整示例 parser.add_argument(--dst_path, typestr, default./frames) parser.add_argument(--new_width, typeint, default340) # 调整输入分辨率 parser.add_argument(--new_height, typeint, default256) parser.add_argument(--frame_rate, typeint, default15) # 控制帧采样率常见问题解决方案FFmpeg路径问题在Python脚本中显式指定FFmpeg路径import os os.environ[PATH] :/usr/local/ffmpeg/bin内存溢出分批次处理大视频文件帧对齐问题保持所有视频的宽高比一致2.2 标签生成与数据集划分标签生成流程中的关键文件说明文件类型格式示例作用classInd.txt1 Action1类别索引映射trainlistXX.txtAction1/v_Action1_g01_c01 1训练集清单testlistXX.txtAction2/v_Action2_g01_c01测试集清单数据集划分建议比例训练集70-80%验证集10-15%测试集10-15%注意划分时应确保每个类别的样本分布均衡避免出现长尾分布问题3. 模型训练与调参策略3.1 配置文件关键参数解析修改dataset_config.py时需要特别注意以下字段# 典型配置示例 event5: { num_class: 5, # 动作类别数 img_path: frames, # 帧存储路径 train_list: trainlist01.txt, val_list: testlist01.txt, modality: RGB, image_tmpl: img_{:05d}.jpg # 必须与帧命名格式一致 }3.2 训练启动与参数优化基础训练命令python main.py event5 RGB \ --arch resnet50 \ --num_segments 8 \ --lr 0.01 \ --lr_steps 20 40 \ --epochs 50 \ --batch-size 32 \ --dropout 0.5 \ --consensus_type avg \ --eval-freq 1 \ --shift \ --shift_div 8参数调优指南batch-size选择8GB显存建议16-3216GB显存建议32-64学习率策略# 阶梯式下降配置示例 --lr 0.01 --lr_steps 20 40 --epochs 50 # 余弦退火配置示例 --lr 0.01 --lr_type cos --epochs 50时序建模参数num_segments视频分段数通常8-16shift_div时序移位比例通常8-164. 实战问题排查与性能优化4.1 常见错误解决方案错误类型可能原因解决方案CUDA内存不足batch-size过大减小batch-size或使用梯度累积标签不匹配文件路径错误检查标签文件中的路径前缀精度波动大学习率过高尝试降低学习率或使用warmup过拟合数据量不足增加数据增强或使用正则化4.2 高级优化技巧数据增强策略# 在transforms.py中添加自定义增强 transforms.Compose([ RandomResizedCrop(224), ColorJitter(brightness0.3, contrast0.3), RandomHorizontalFlip(), ])混合精度训练# 启动命令添加--fp16参数 python main.py ... --fp16模型微调技巧冻结底层卷积层前3-5个block对全连接层使用更大学习率使用标签平滑label smoothing在实际项目中我们发现将num_segments设置为16配合余弦退火学习率调度可以在UCF101格式的数据集上获得最佳性能平衡。对于长视频处理建议先进行关键帧提取再输入模型可以显著提升处理效率。