TSM行为识别实战：从UCF101抽帧到模型训练，保姆级避坑指南（附完整代码）

TSM行为识别实战从UCF101抽帧到模型训练全流程解析行为识别作为计算机视觉领域的重要研究方向正在智能监控、人机交互、体育分析等场景中发挥越来越大的作用。本文将手把手带你完成基于TSMTemporal Shift Module模型的行为识别全流程实战从UCF101数据集处理到模型训练调优每个步骤都配有详细的操作指南和避坑建议。1. 环境准备与数据获取在开始实战之前我们需要准备好开发环境和数据集。以下是推荐的环境配置硬件配置GPUNVIDIA GTX 1080 Ti或更高建议显存≥8GBCPU4核以上内存16GB以上存储空间至少100GB可用空间UCF101原始视频抽帧图片占用较大空间软件环境操作系统Ubuntu 18.04/20.04推荐或Windows 10Python3.7或3.8PyTorch1.7.0或更高版本CUDA10.2或11.1需与PyTorch版本匹配cuDNN与CUDA对应的版本提示建议使用conda创建虚拟环境避免依赖冲突。可使用以下命令创建环境conda create -n tsm python3.8 conda activate tsm pip install torch torchvision torchaudioUCF101数据集下载与结构UCF101是行为识别领域广泛使用的基准数据集包含101类动作的13320个视频片段。数据集可从官网下载解压后目录结构如下UCF101/ ├── ApplyEyeMakeup/ │ ├── v_ApplyEyeMakeup_g01_c01.avi │ ├── v_ApplyEyeMakeup_g01_c02.avi │ └── ... ├── ApplyLipstick/ │ ├── v_ApplyLipstick_g01_c01.avi │ └── ... └── ...2. 视频抽帧与数据预处理TSM模型需要以视频帧序列作为输入因此我们需要将视频文件转换为图像序列。以下是详细的抽帧流程2.1 安装FFmpegFFmpeg是处理视频的必备工具可通过以下命令安装# Ubuntu sudo apt-get update sudo apt-get install ffmpeg # Windows # 从官网下载编译好的二进制文件并添加到PATH环境变量2.2 编写抽帧脚本创建vid2img_ucf101.py脚本实现批量视频抽帧import os import subprocess import argparse def video_to_frames(video_path, output_dir, frame_size256): 将单个视频转换为帧序列 if not os.path.exists(output_dir): os.makedirs(output_dir) cmd fffmpeg -i {video_path} -vf scale{frame_size}:-1 -q:v 2 {output_dir}/image_%05d.jpg subprocess.call(cmd, shellTrue) def process_class(class_dir, output_root, class_name): 处理一个类别下的所有视频 class_output_dir os.path.join(output_root, class_name) if not os.path.exists(class_output_dir): os.makedirs(class_output_dir) for video_file in os.listdir(class_dir): if not video_file.endswith(.avi): continue video_name os.path.splitext(video_file)[0] video_path os.path.join(class_dir, video_file) output_dir os.path.join(class_output_dir, video_name) video_to_frames(video_path, output_dir) if __name__ __main__: parser argparse.ArgumentParser() parser.add_argument(--input_dir, requiredTrue, helpUCF101视频根目录) parser.add_argument(--output_dir, requiredTrue, help输出图片根目录) args parser.parse_args() for class_name in os.listdir(args.input_dir): class_dir os.path.join(args.input_dir, class_name) if os.path.isdir(class_dir): process_class(class_dir, args.output_dir, class_name)运行脚本python vid2img_ucf101.py --input_dir /path/to/UCF101 --output_dir /path/to/output_frames2.3 生成标签文件创建gen_label_ucf101.py脚本生成训练所需的标签文件import os import glob def generate_label_files(frame_root, label_dir, class_ind_file, train_list_file, test_list_file): # 读取类别索引 class_ind {} with open(class_ind_file) as f: for line in f: idx, name line.strip().split() class_ind[name] int(idx) - 1 # 转换为0-based索引 # 生成训练集标签 train_samples [] with open(train_list_file) as f: for line in f: parts line.strip().split(/) class_name parts[0] video_name parts[1].split(.)[0] label class_ind[class_name] frame_dir os.path.join(frame_root, class_name, video_name) frame_count len(glob.glob(os.path.join(frame_dir, *.jpg))) if frame_count 0: # 确保有有效帧 train_samples.append(f{frame_dir} {frame_count} {label}\n) # 生成测试集标签 test_samples [] with open(test_list_file) as f: for line in f: parts line.strip().split(/) class_name parts[0] video_name parts[1].split(.)[0] label class_ind[class_name] frame_dir os.path.join(frame_root, class_name, video_name) frame_count len(glob.glob(os.path.join(frame_dir, *.jpg))) if frame_count 0: test_samples.append(f{frame_dir} {frame_count} {label}\n) # 写入文件 os.makedirs(label_dir, exist_okTrue) with open(os.path.join(label_dir, ucf101_rgb_train_split_1.txt), w) as f: f.writelines(train_samples) with open(os.path.join(label_dir, ucf101_rgb_val_split_1.txt), w) as f: f.writelines(test_samples) if __name__ __main__: generate_label_files( frame_root/path/to/output_frames, label_dir/path/to/labels, class_ind_file/path/to/classInd.txt, train_list_file/path/to/trainlist01.txt, test_list_file/path/to/testlist01.txt )3. TSM模型训练配置3.1 获取TSM官方代码克隆TSM官方仓库并安装依赖git clone https://github.com/mit-han-lab/temporal-shift-module.git cd temporal-shift-module pip install -r requirements.txt3.2 修改配置文件在ops/dataset_config.py中修改数据集路径ucf101 { rgb_train: /path/to/ucf101_rgb_train_split_1.txt, rgb_val: /path/to/ucf101_rgb_val_split_1.txt, root_dataset: /path/to/output_frames, num_class: 101, }3.3 下载预训练权重TSM提供了在Kinetics数据集上预训练的权重可从以下链接下载ResNet50版本TSM_kinetics_RGB_resnet50_shift8_blockres_avg_segment8_e50.pthMobileNetV2版本TSM_kinetics_RGB_mobilenetv2_shift8_blockres_avg_segment8_e100_dense.pth将下载的权重文件放入pretrained/目录。4. 模型训练与调优4.1 训练命令解析TSM提供了多种训练选项以下是常用参数说明参数说明推荐值--arch模型架构resnet50或mobilenetv2--num_segment时间片段数8--batch-size批次大小16-32根据GPU显存调整--lr初始学习率0.001--lr_steps学习率衰减步数10 20--epochs训练轮数25-50--dropoutDropout率0.5-0.8--shift_div时间移位除数8--tune_from预训练权重路径pretrained/xxx.pthResNet50训练命令python main.py ucf101 RGB \ --arch resnet50 \ --num_segment 8 \ --batch-size 16 \ --lr 0.001 \ --lr_steps 10 20 \ --epochs 25 \ --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.pthMobileNetV2训练命令python main.py ucf101 RGB \ --arch mobilenetv2 \ --num_segment 8 \ --batch-size 32 \ --lr 0.001 \ --lr_steps 10 20 \ --epochs 25 \ --dropout 0.5 \ --consensus_type avg \ --eval-freq 1 \ --shift \ --shift_div 8 \ --shift_place blockres \ --tune_from pretrained/TSM_kinetics_RGB_mobilenetv2_shift8_blockres_avg_segment8_e100_dense.pth4.2 常见问题解决问题1预训练权重加载失败MobileNetV2版本可能出现权重加载错误解决方法是在main.py中修改权重加载逻辑if args.arch mobilenetv2: model_dict model.module.state_dict() # 处理权重名称不匹配问题 new_sd {} for k, v in sd.items(): if k.startswith(module.): new_k k[7:] # 去除module.前缀 new_sd[new_k] v else: new_sd[k] v model_dict.update(new_sd) model.module.load_state_dict(model_dict)问题2显存不足如果遇到CUDA out of memory错误可以尝试减小batch-size如从32降到16减小输入分辨率修改--new_length和--new_width参数使用梯度累积添加--accumulate_grad参数问题3训练精度不理想可以尝试以下调优策略增加数据增强修改transform.py中的增强策略调整学习率策略如使用余弦退火增加训练轮数--epochs尝试不同的shift_div值如4或165. 模型测试与部署5.1 视频行为识别Demo创建video_demo.py实现视频行为识别import torch import cv2 import numpy as np from models import TSN from transforms import GroupScale, GroupCenterCrop, Stack, ToTorchFormatTensor, GroupNormalize # 初始化模型 num_class 101 arch resnet50 pretrain checkpoints/TSM_ucf101_RGB_resnet50_shift8_blockres_avg_segment8_e25.pth model TSN(num_class, 8, RGB, base_modelarch, consensus_typeavg, dropout0.8, shift_div8) checkpoint torch.load(pretrain) model.load_state_dict(checkpoint[state_dict]) model.eval() # 定义预处理 transform torchvision.transforms.Compose([ GroupScale(256), GroupCenterCrop(224), Stack(), ToTorchFormatTensor(), GroupNormalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) # 加载类别标签 with open(ucf101_labels.txt) as f: class_names [line.strip() for line in f] # 视频处理函数 def process_video(video_path): cap cv2.VideoCapture(video_path) frames [] while True: ret, frame cap.read() if not ret: break # 转换为RGB frame cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) frames.append(Image.fromarray(frame)) # 采样8帧 num_frames len(frames) indices np.linspace(0, num_frames-1, 8, dtypenp.int32) sampled_frames [frames[i] for i in indices] # 预处理 input_tensor transform(sampled_frames).unsqueeze(0) # 推理 with torch.no_grad(): output model(input_tensor) # 获取预测结果 pred output.argmax(dim1).item() return class_names[pred] # 测试视频 video_path test_video.avi prediction process_video(video_path) print(fPredicted action: {prediction})5.2 性能优化技巧模型量化使用PyTorch的量化功能减小模型大小提升推理速度model torch.quantization.quantize_dynamic( model, {torch.nn.Linear}, dtypetorch.qint8 )多线程预处理使用DataLoader的num_workers参数加速数据加载TensorRT加速将模型转换为TensorRT引擎提升推理性能帧采样策略优化根据视频长度动态调整采样间隔确保覆盖关键动作在实际项目中TSM模型在UCF101数据集上通常能达到85%-90%的top-1准确率。通过合理的调参和数据增强可以进一步提升模型性能。