ST-GCN实战指南：Ubuntu20.04+GPU环境下的复现与优化

1. 环境准备搭建Ubuntu20.04GPU基础环境在开始ST-GCN的复现之前我们需要先搭建一个稳定的基础环境。Ubuntu20.04是目前比较稳定的Linux发行版对NVIDIA显卡的支持也比较好。我使用的是RTX 3070显卡搭配CUDA 11.1和cuDNN 8.1.0这个组合在实际测试中表现很稳定。首先需要确认你的显卡驱动已经正确安装。可以通过以下命令检查nvidia-smi这个命令会显示显卡的基本信息和CUDA版本。需要注意的是这里显示的CUDA版本是显卡驱动支持的最高版本而不是你实际安装的CUDA工具包版本。接下来安装CUDA 11.1。建议使用runfile方式安装这样可以更灵活地选择组件wget https://developer.download.nvidia.com/compute/cuda/11.1.0/local_installers/cuda_11.1.0_455.23.05_linux.run sudo sh cuda_11.1.0_455.23.05_linux.run安装过程中记得取消勾选显卡驱动选项如果你已经安装了合适的驱动。安装完成后需要将CUDA添加到环境变量中echo export PATH/usr/local/cuda-11.1/bin:$PATH ~/.bashrc echo export LD_LIBRARY_PATH/usr/local/cuda-11.1/lib64:$LD_LIBRARY_PATH ~/.bashrc source ~/.bashrccuDNN的安装稍微麻烦一些需要先到NVIDIA官网下载对应版本的库文件然后手动解压并复制到CUDA目录tar -xzvf cudnn-11.1-linux-x64-v8.1.0.77.tgz sudo cp cuda/include/cudnn*.h /usr/local/cuda-11.1/include sudo cp cuda/lib64/libcudnn* /usr/local/cuda-11.1/lib64 sudo chmod ar /usr/local/cuda-11.1/include/cudnn*.h /usr/local/cuda-11.1/lib64/libcudnn*最后验证CUDA和cuDNN是否安装成功nvcc --version cat /usr/local/cuda/include/cudnn_version.h | grep CUDNN_MAJOR -A 22. 源码获取与依赖安装ST-GCN的官方实现使用PyTorch框架我们需要先准备好Python环境。我推荐使用conda来管理Python环境这样可以避免与系统Python产生冲突。首先创建并激活conda环境conda create -n stgcn python3.6 conda activate stgcn然后安装PyTorch 1.8.0这个版本与CUDA 11.1兼容性最好conda install pytorch1.8.0 torchvision0.9.0 torchaudio0.8.0 cudatoolkit11.1 -c pytorch -c conda-forge接下来获取ST-GCN的源码。国内用户可以使用gitee镜像速度会快很多git clone https://gitee.com/chenhongqiong/st-gcn.git cd st-gcn安装其他依赖项conda install ffmpeg pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple这里有个小技巧如果遇到某些包安装失败可以尝试先用conda安装基础版本再用pip升级到指定版本。我在安装scipy时就遇到了这个问题解决方法如下conda install scipy pip install --upgrade scipy1.5.43. OpenPose的编译与安装ST-GCN依赖OpenPose来提取人体姿态特征这是整个复现过程中最复杂的部分之一。OpenPose的编译过程可能会遇到各种问题下面我会详细介绍如何避坑。首先安装系统依赖sudo apt-get update sudo apt-get install cmake git unzip libprotobuf-dev libleveldb-dev liblmdb-dev sudo apt-get install libsnappy-dev libhdf5-serial-dev protobuf-compiler sudo apt-get install --no-install-recommends libboost-all-dev sudo apt-get install libatlas-base-dev libopenblas-dev sudo apt-get install libopencv-dev获取OpenPose源码建议使用1.5.0版本git clone https://github.com/CMU-Perceptual-Computing-Lab/openpose.git cd openpose git checkout v1.5.0 git submodule update --init --recursive --remote编译前的准备工作mkdir build cd build cmake-gui ..在CMake图形界面中需要特别注意以下几个配置勾选BUILD_PYTHON选项设置Python路径为你conda环境中的Python设置CUDA_ARCH_BIN为你的显卡计算能力RTX 3070是8.6配置完成后点击Generate生成Makefile然后开始编译make -jnproc这个过程可能会持续1-2小时取决于你的机器性能。如果遇到编译错误最常见的原因是内存不足可以尝试减少并行编译线程数比如使用make -j4。编译完成后还需要下载OpenPose的预训练模型。由于网络原因可能需要手动下载这些模型文件并放到正确的目录cd models wget http://posefs1.perception.cs.cmu.edu/OpenPose/models/pose/coco/pose_iter_440000.caffemodel wget http://posefs1.perception.cs.cmu.edu/OpenPose/models/face/pose_iter_116000.caffemodel wget http://posefs1.perception.cs.cmu.edu/OpenPose/models/hand/pose_iter_102000.caffemodel最后测试OpenPose是否安装成功./examples/openpose/openpose.bin --image_dir examples/media/ --face --hand如果能看到处理后的图片输出说明安装成功。4. 数据集准备与预处理ST-GCN支持多种数据集最常用的是NTU RGBD和Kinetics-skeleton。这里我以NTU RGBD为例介绍数据准备过程。首先需要申请NTU RGBD数据集。访问ROSE Lab官网填写申请表通常1-2天内会收到下载链接。下载完成后你会得到一个约5.8GB的压缩文件。解压数据集unzip nturgbd_skeletons.zipST-GCN需要特定的数据格式我们需要运行预处理脚本python tools/ntu_gendata.py --data_path /path/to/nturgbd_skeletons这个脚本会生成四个文件nturgbd_skeletons_train_xview.npynturgbd_skeletons_val_xview.npynturgbd_skeletons_train_xsub.npynturgbd_skeletons_val_xsub.npy对于Kinetics-skeleton数据集处理方式类似python tools/kinetics_gendata.py --data_path /path/to/kinetics-skeleton在实际操作中我发现数据集路径设置是个常见的错误源。建议使用绝对路径并在config文件中仔细检查路径配置。如果遇到can not find pose estimation results错误通常就是路径配置不正确导致的。5. 模型训练与调优环境准备好后就可以开始训练模型了。ST-GCN提供了多个配置文件对应不同的数据集和实验设置。以NTU RGBD cross-view实验为例首先修改配置文件config/st_gcn/ntu-xview/train.yamlwork_dir: ./work_dir device: [0] # 使用单GPU训练 batch_size: 32 test_batch_size: 32开始训练python main.py recognition -c config/st_gcn/ntu-xview/train.yaml训练过程中可能会遇到几个常见问题GPU内存不足可以尝试减小batch_size或者使用梯度累积技术训练速度慢检查CUDA和cuDNN是否正确安装nvidia-smi查看GPU利用率损失不下降尝试调整学习率或者检查数据预处理是否正确我在RTX 3070上的训练参数如下batch_size: 56base_lr: 0.1使用StepLR调度器每20个epoch学习率乘以0.1训练过程中可以监控GPU使用情况watch -n 1 nvidia-smi如果训练中断可以修改train.yaml中的start_epoch和weights参数继续训练start_epoch: 20 weights: ./work_dir/recognition/ntu-xview/ST_GCN/epoch20_model.pt6. 模型测试与性能分析训练完成后我们需要评估模型性能。使用以下命令测试模型python main.py recognition -c config/st_gcn/ntu-xview/test.yaml --weights ./work_dir/recognition/ntu-xview/ST_GCN/epoch80_model.pt在我的测试中NTU RGBD cross-view的准确率达到了87.47%比论文报告的结果低了约1.4个百分点。这可能是因为训练epoch不够多论文训练了100个epoch数据增强策略有差异随机种子不同导致的波动为了进一步提升性能可以尝试以下优化策略增加训练epoch到100-120使用更复杂的数据增强调整网络结构如增加通道数尝试不同的优化器参数对于Kinetics-skeleton数据集测试命令类似python main.py recognition -c config/st_gcn/kinetics-skeleton/test.yaml --weights ./path/to/model.pt7. 常见问题解决方案在实际复现过程中我遇到了不少问题这里总结几个典型的案例问题1OpenPose编译失败错误信息CMake Error at cmake/Cuda.cmake:227 (message): cuDNN version 3 is required. 解决方案修改OpenPpose/cmake/Cuda.cmake文件将cudnn.h改为cudnn_version.h问题2ST-GCN运行时维度不匹配错误信息RuntimeError: Given groups1, weight of size [60, 256, 1, 1], expected input[8, 3, 1, 1] to have 256 channels 解决方案检查数据预处理是否正确确保输入数据的维度与模型匹配问题3双卡训练报错解决方案目前ST-GCN的官方实现似乎对多GPU支持不够完善建议使用单卡训练。如果必须使用多卡可以尝试修改net/st_gcn.py中的forward函数问题4Demo运行失败解决方案确保OpenPose路径配置正确并且视频文件路径无误。可以先用OpenPose自带的demo测试环境是否正常问题5显存不足解决方案减小batch_size或者使用更小的模型。RTX 3070(8GB)建议batch_size不超过648. 实际应用与扩展成功复现ST-GCN后我们可以尝试将其应用到实际场景中。以下是一些可能的扩展方向实时动作识别修改demo_old.py接入摄像头实时输入自定义动作识别在自己的数据集上微调模型模型轻量化尝试知识蒸馏或量化技术降低模型计算量多模态融合结合RGB视频流和骨骼数据提升准确率我在尝试实时动作识别时发现性能瓶颈主要在OpenPose的姿态估计部分。一个实用的优化策略是降低OpenPose的输入分辨率或者使用更轻量级的姿态估计模型。对于想尝试自定义数据集的开发者需要注意数据格式需要与NTU RGBD保持一致。具体来说每个样本应该包含18个关节点的3D坐标时间序列数据通常30帧/秒对应的动作标签最后提醒一点ST-GCN对输入数据的质量比较敏感。在实际应用中可能需要加入一些后处理步骤比如关节点的平滑滤波缺失数据的插值异常值的检测与修正