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

weixin_30882895

638人浏览 · 2026-04-15 17:57:43

weixin_30882895 · 2026-04-15 17:57:43 发布

1. 环境准备：搭建Ubuntu20.04+GPU基础环境

在开始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 ~/.bashrc

cuDNN的安装稍微麻烦一些，需要先到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 a+r /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 2

2. 源码获取与依赖安装

ST-GCN的官方实现使用PyTorch框架，我们需要先准备好Python环境。我推荐使用conda来管理Python环境，这样可以避免与系统Python产生冲突。

首先创建并激活conda环境：

conda create -n stgcn python=3.6
conda activate stgcn

然后安装PyTorch 1.8.0，这个版本与CUDA 11.1兼容性最好：

conda install pytorch==1.8.0 torchvision==0.9.0 torchaudio==0.8.0 cudatoolkit=11.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 scipy==1.5.4

3. 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 -j`nproc`

这个过程可能会持续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 RGB+D和Kinetics-skeleton。这里我以NTU RGB+D为例介绍数据准备过程。

首先需要申请NTU RGB+D数据集。访问ROSE Lab官网填写申请表，通常1-2天内会收到下载链接。下载完成后你会得到一个约5.8GB的压缩文件。

解压数据集：

unzip nturgb+d_skeletons.zip

ST-GCN需要特定的数据格式，我们需要运行预处理脚本：

python tools/ntu_gendata.py --data_path /path/to/nturgb+d_skeletons

这个脚本会生成四个文件：

nturgb+d_skeletons_train_xview.npy
nturgb+d_skeletons_val_xview.npy
nturgb+d_skeletons_train_xsub.npy
nturgb+d_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 RGB+D cross-view实验为例，首先修改配置文件config/st_gcn/ntu-xview/train.yaml：

work_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: 56
base_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.pt

6. 模型测试与性能分析

训练完成后，我们需要评估模型性能。使用以下命令测试模型：

python main.py recognition -c config/st_gcn/ntu-xview/test.yaml --weights ./work_dir/recognition/ntu-xview/ST_GCN/epoch80_model.pt

在我的测试中，NTU RGB+D 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.pt

7. 常见问题解决方案

在实际复现过程中，我遇到了不少问题，这里总结几个典型的案例：

问题1：OpenPose编译失败 错误信息：CMake Error at cmake/Cuda.cmake:227 (message): cuDNN version >3 is required. 解决方案：修改OpenPpose/cmake/Cuda.cmake文件，将cudnn.h改为cudnn_version.h

问题2：ST-GCN运行时维度不匹配 错误信息：RuntimeError: Given groups=1, 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函数

问题4：Demo运行失败 解决方案：确保OpenPose路径配置正确，并且视频文件路径无误。可以先用OpenPose自带的demo测试环境是否正常

问题5：显存不足 解决方案：减小batch_size，或者使用更小的模型。RTX 3070(8GB)建议batch_size不超过64