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

1. UCF101数据集准备实战指南

第一次接触行为识别任务时，最让人头疼的就是数据准备环节。UCF101作为行为识别领域的经典数据集，包含101类人类动作视频，总计13,320个视频片段。但原始视频文件需要经过特定处理才能用于TSM模型训练。下面分享我踩坑后总结的高效处理方法。

首先需要明确UCF101的标准目录结构。下载解压后会得到UCF-101主文件夹，内含101个子文件夹，每个子文件夹对应一个动作类别（如ApplyEyeMakeup、BasketballDunk等），子文件夹内是该类别的所有avi格式视频文件。建议先创建如下目录结构：

UCF101/ ├── videos/ # 存放原始视频 ├── frames/ # 存放抽帧结果 ├── annotations/ # 存放标签文件 └── splits/ # 存放训练测试划分

视频抽帧是关键步骤，这里推荐使用FFmpeg工具。我编写了一个自动化脚本，可以批量处理所有视频：

import os import subprocess def extract_frames(video_path, output_dir, fps=30): if not os.path.exists(output_dir): os.makedirs(output_dir) command = f"ffmpeg -i {video_path} -r {fps} {output_dir}/frame_%05d.jpg" subprocess.call(command, shell=True) # 批量处理示例 for class_name in os.listdir("UCF101/videos"): class_path = os.path.join("UCF101/videos", class_name) for video_file in os.listdir(class_path): video_path = os.path.join(class_path, video_file) output_dir = os.path.join("UCF101/frames", class_name, video_file[:-4]) extract_frames(video_path, output_dir)

注意：实际运行时需要根据硬件配置调整fps参数。在GTX 1080Ti上处理完整数据集约需2小时，建议先抽取部分类别测试。

标签文件准备是另一个易错点。UCF101官方提供三个标准训练测试划分（split1/2/3），需要将视频路径映射到对应类别ID。我建议使用以下Python代码生成TSM所需的标签格式：

def generate_label_files(split_file, output_file): with open(split_file) as f: lines = f.readlines() with open(output_file, 'w') as f_out: for line in lines: video_path, class_id = line.strip().split() frame_dir = os.path.join("UCF101/frames", video_path.split('/')[0], video_path.split('/')[1][:-4]) frame_count = len(os.listdir(frame_dir)) f_out.write(f"{frame_dir} {frame_count} {class_id}\n") # 为三个划分生成标签 for i in [1,2,3]: generate_label_files(f"UCF101/splits/trainlist0{i}.txt", f"UCF101/annotations/train_split{i}.txt") generate_label_files(f"UCF101/splits/testlist0{i}.txt", f"UCF101/annotations/val_split{i}.txt")

2. TSM模型配置详解

拿到处理好的数据后，模型配置是下一个关键环节。TSM（Temporal Shift Module）作为高效的时序建模方法，其配置文件需要特别注意几个核心参数。

首先是数据集路径配置，需要修改ops/dataset_config.py文件：

ucf101 = { 'root_dataset': './UCF101/frames/', # 抽帧图片根目录 'train_source': './UCF101/annotations/train_split1.txt', # 训练集标签 'val_source': './UCF101/annotations/val_split1.txt', # 验证集标签 'num_class': 101, # 类别数（若使用子集需调整） }

模型架构选择方面，TSM支持多种backbone。对于UCF101数据集，我推荐以下两种配置方案：

1. ResNet50方案：

   • 优势：精度高（top1准确率约94%）
   • 缺点：计算量较大（约16G显存）

2. MobileNetV2方案：

   • 优势：轻量化（仅需4G显存）
   • 缺点：精度稍低（约89%）

训练参数配置直接影响模型性能。以下是经过验证的最佳参数组合：

对于初次尝试的建议：

• 使用MobileNetV2版本快速验证流程
• 8个时间片段（num_segment）平衡了精度和效率
• 学习率采用阶梯式衰减策略
• 启用--shift_div=8参数激活时序移位功能

3. 训练过程全解析

配置完成后，真正的挑战才开始。以下是我在四块2080Ti显卡上的实战经验。

启动训练的命令示例如下：

# ResNet50版本 python main.py ucf101 RGB \ --arch resnet50 \ --num_segment 8 \ --gd 20 \ --lr 0.001 \ --lr_steps 10 20 \ --epochs 25 \ --batch-size 16 \ -j 16 \ --dropout 0.8 \ --consensus_type avg \ --eval-freq 1 \ --shift \ --shift_div 8 \ --shift_place blockres \ --tune_from pretrained/TSM_kinetics_RGB_resnet50_shift8_blockres_avg_segment8_e50.pth

训练过程中需要特别关注以下几个指标：

1. Loss曲线：

   • 正常情况：训练loss应平稳下降，验证loss同步下降
   • 异常情况：两者差距过大可能过拟合

2. 准确率变化：

   • 初期应快速上升（前5个epoch）
   • 后期缓慢收敛（最后5个epoch提升<1%）

3. 显存占用：

   • ResNet50：约15GB/GPU
   • MobileNetV2：约7GB/GPU

常见问题及解决方案：

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

• 现象：报错"Missing keys in state_dict"
• 解决：修改model.py中的权重加载逻辑：

if args.arch == "mobilenetv2": sd = {k.replace('base_model.', ''): v for k,v in sd.items()}

问题2：视频帧数不一致

• 现象：RuntimeError: inconsistent frame numbers
• 解决：确保所有视频至少包含num_segment*16帧（对于8片段需128帧）

问题3：验证准确率波动大

• 现象：val_acc忽高忽低
• 解决：增大batch_size或减小学习率

训练完成后，模型会保存在checkpoint目录。建议使用最后5个epoch的平均权重作为最终模型，可以通过以下命令实现：

python tools/average_weights.py \ --input checkpoints/ucf101_resnet50_shift8_blockres_avg_segment8_e25 \ --num_epoch 5 \ --output final_model.pth

4. 模型测试与应用

训练出的模型需要经过严格测试才能投入实用。TSM官方代码缺少现成的测试脚本，我开发了一套完整的测试流程。

首先准备测试视频处理脚本：

import torch from model import TSN from transforms import GroupScale, GroupCenterCrop, Stack, ToTorchFormatTensor def load_model(checkpoint_path, num_class): model = TSN(num_class, num_segments=8, modality='RGB', base_model='resnet50', consensus_type='avg', dropout=0.8) checkpoint = torch.load(checkpoint_path) model.load_state_dict(checkpoint['state_dict']) return model.eval() def preprocess_frame(frame): transform = torchvision.transforms.Compose([ GroupScale(256), GroupCenterCrop(224), Stack(), ToTorchFormatTensor(), ]) return transform([Image.fromarray(frame)]).unsqueeze(0)

视频预测核心逻辑：

def predict_video(model, video_path): cap = cv2.VideoCapture(video_path) buffers = [] while True: ret, frame = cap.read() if not ret: break inputs = preprocess_frame(frame) with torch.no_grad(): outputs = model(inputs) buffers.append(outputs) if len(buffers) == 8: # 与num_segment一致 final_output = torch.mean(torch.stack(buffers), dim=0) pred_class = torch.argmax(final_output).item() buffers = [] # 清空缓存 return pred_class

对于实际部署，我推荐两种方案：

方案A：在线推理（低延迟）

• 使用TensorRT加速
• 优化后的ResNet50在T4显卡上可达80FPS
• 适合实时监控场景

方案B：批量处理（高吞吐）

• 使用多进程并行
• 单卡可同时处理16-32路视频
• 适合视频分析场景

性能优化技巧：

1. 启用半精度推理（FP16）可提升40%速度
2. 使用多尺度裁剪（3-crop）可提升2-3%准确率
3. 时序分段重叠采样（overlap=0.5）提升时序建模能力

最后分享一个实用技巧：当遇到识别不准的情况时，可以尝试以下调整：

• 增加num_segment到16（需更多显存）
• 使用更大的输入分辨率（从224x224到320x320）
• 在Kinetics数据集上预训练后再微调

经过完整流程训练出的TSM模型，在UCF101测试集上可以达到以下性能：

实际项目中，我通常会先用MobileNetV2快速原型验证，再根据需要切换到ResNet50追求更高精度。记住，模型选择最终要服务于业务需求，在精度和效率之间找到最佳平衡点才是工程实践的精髓。