当前位置：首页 > news >正文

TSM行为识别实战：从UCF101数据集准备到模型训练，保姆级避坑指南

news 2026/4/11 9:21:07

TSM行为识别实战：从UCF101数据集准备到模型训练全流程解析

在计算机视觉领域，行为识别一直是一个极具挑战性的研究方向。不同于静态图像分类，视频行为识别需要模型能够理解时间维度上的动作变化。TSM（Temporal Shift Module）作为近年来备受关注的时间建模方法，通过巧妙地在2D CNN中引入时序建模能力，在保持高效推理的同时显著提升了行为识别性能。本文将手把手带你完成从UCF101数据集准备到TSM模型训练的全过程，特别针对实际项目中容易遇到的"坑点"提供解决方案。

1. UCF101数据集准备与预处理

UCF101作为行为识别领域的基准数据集，包含101类动作的13320个视频片段。但在实际应用中，我们往往不需要使用全部类别。以下是经过优化的数据处理流程：

1.1 数据集下载与结构分析

首先从官网获取UCF101数据集，其目录结构通常如下：

UCF101/ ├── ApplyEyeMakeup/ │ ├── v_ApplyEyeMakeup_g01_c01.avi │ └── ... ├── ApplyLipstick/ │ ├── v_ApplyLipstick_g01_c01.avi │ └── ... └── ...

关键注意事项：

视频文件以.avi格式存储
每个子目录对应一个动作类别
文件名包含视频属性信息（如g01表示组别，c01表示摄像机角度）

1.2 视频抽帧处理

TSM等行为识别模型通常以视频帧序列作为输入。我们使用FFmpeg进行抽帧：

import os import subprocess def extract_frames(video_path, output_dir, fps=25, size=(340, 256)): """使用FFmpeg从视频中提取帧""" if not os.path.exists(output_dir): os.makedirs(output_dir) cmd = f"ffmpeg -i {video_path} -r {fps} -s {size[0]}x{size[1]} -q:v 2 {output_dir}/image_%05d.jpg" subprocess.call(cmd, shell=True)

常见问题解决方案：

抽帧速度慢：降低目标分辨率或帧率
内存不足：分批处理视频文件
时间戳错乱：确保使用%05d格式保证文件名排序正确

1.3 生成标签文件

UCF101官方提供了训练/测试划分文件。我们需要将其转换为模型可读的格式：

def generate_label_file(video_dir, output_path): classes = sorted(os.listdir(video_dir)) class_to_idx = {cls_name: i for i, cls_name in enumerate(classes)} with open(output_path, 'w') as f: for cls_name in classes: cls_dir = os.path.join(video_dir, cls_name) for video in os.listdir(cls_dir): frame_dir = os.path.join(cls_dir, video.split('.')[0]) frame_count = len(os.listdir(frame_dir)) f.write(f"{frame_dir} {frame_count} {class_to_idx[cls_name]}\n")

提示：对于大型数据集，建议使用多进程加速标签生成过程

2. TSM模型环境配置

2.1 依赖安装

推荐使用conda创建Python 3.7环境：

conda create -n tsm python=3.7 conda activate tsm pip install torch==1.8.0 torchvision==0.9.0 pip install opencv-python ffmpeg-python

2.2 代码库克隆与修改

从官方仓库克隆TSM代码：

git clone https://github.com/mit-han-lab/temporal-shift-module.git cd temporal-shift-module

关键修改点：

在dataset_config.py中更新数据集路径
根据GPU内存调整batch_size参数
修改main.py中的学习率调度策略

3. 模型训练与调优

3.1 训练命令解析

基础训练命令示例：

python main.py ucf101 RGB \ --arch resnet50 \ --num_segments 8 \ --lr 0.001 \ --epochs 50 \ --batch-size 32 \ --dropout 0.5 \ --consensus_type=avg \ --shift --shift_div=8 --shift_place=blockres \ --tune_from=pretrained/TSM_kinetics_RGB_resnet50_shift8_blockres_avg_segment8_e50.pth

参数说明：

num_segments：输入视频划分的片段数
shift_div：控制时序移位比例
tune_from：预训练权重路径

3.2 常见训练问题解决

问题1：预训练权重加载失败

解决方案：修改模型加载逻辑

# 在main.py中找到权重加载部分，添加以下处理 if 'module.' in list(sd.keys())[0]: sd = {k.replace('module.', ''): v for k,v in sd.items()}

问题2：GPU内存不足

优化策略：

减小batch_size
使用梯度累积
尝试更小的模型架构（如MobileNetV2）

问题3：过拟合

应对方法：

增加dropout值（0.5-0.8）
添加数据增强（随机裁剪、水平翻转等）
使用早停策略

4. 模型测试与部署

4.1 视频推理实现

官方代码缺少直接的视频输入接口，我们需要自行实现：

import torch import cv2 import numpy as np from models import TSN def video_inference(model, video_path, num_segments=8): # 初始化模型 model.eval() # 视频读取 cap = cv2.VideoCapture(video_path) frames = [] while cap.isOpened(): ret, frame = cap.read() if not ret: break frame = preprocess(frame) # 预处理函数 frames.append(frame) # 分段采样 segment_length = len(frames) // num_segments indices = [i*segment_length for i in range(num_segments)] inputs = torch.stack([frames[i] for i in indices]) # 推理 with torch.no_grad(): outputs = model(inputs.unsqueeze(0)) return outputs.argmax().item()