STGCN（时空图卷积网络）详解：原理、结构与代码实现

创作时间:

作者:

@小白创作中心

STGCN（时空图卷积网络）详解：原理、结构与代码实现

引用

CSDN

https://blog.csdn.net/qq_42980908/article/details/143606847

时空图卷积网络（STGCN）是一种结合了图卷积神经网络（GCN）和时间卷积网络（TCN）的深度学习模型，主要用于基于骨架图的动作识别任务。本文将详细介绍STGCN的原理、网络结构及其PyTorch实现。

STGCN介绍

STGCN是基于图卷积（GCN）的扩展，用于处理时空图数据。其输入是一系列骨架图，其中每个节点对应人体的一个关节，边分为两类：

空间边：符合节点自然连通性（图1中淡蓝色线条）
时间边：跨越连续时间步长连接相同节点（淡绿色线条）

网络结构

STGCN的网络结构主要包括三个部分：

归一化：对输入数据进行归一化处理
时空变化：通过多个ST-GCN块，每个块中交替使用GCN和TCN
输出：使用平均池化和全连接层对特征进行分类

ST-GCN块的详细结构

每个ST-GCN块包含以下步骤：

引入一个可学习的权重矩阵，与邻接矩阵按位相乘，得到加权后的邻接矩阵
将加权后的邻接矩阵与输入数据送入GCN进行运算
利用TCN网络实现时间维度信息的聚合

代码实现

以下是STGCN的完整PyTorch实现代码：

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

# 空域图卷积
class SpatialGraphConvolution(nn.Module):
    def __init__(self, in_channels, out_channels, s_kernel_size):
        super().__init__()
        self.s_kernel_size = s_kernel_size
        self.conv = nn.Conv2d(in_channels=in_channels,
                              out_channels=out_channels * s_kernel_size,
                              kernel_size=1)
    def forward(self, x, A):
        x = self.conv(x)
        n, kc, t, v = x.size()
        x = x.view(n, self.s_kernel_size, kc//self.s_kernel_size, t, v)
        x = torch.einsum('nkctv,kvw->nctw', (x, A))
        return x.contiguous()

# STGCN块
class STGCN_block(nn.Module):
    def __init__(self, in_channels, out_channels, stride, t_kernel_size, A_size, dropout=0.5):
        super().__init__()
        self.s_gcn = SpatialGraphConvolution(in_channels=in_channels,
                                            out_channels=out_channels,
                                            s_kernel_size=A_size[0])
        self.M = nn.Parameter(torch.ones(A_size))
        self.t_gcn = nn.Sequential(nn.BatchNorm2d(out_channels),
                                  nn.ReLU(),
                                  nn.Dropout(dropout),
                                  nn.Conv2d(out_channels,
                                            out_channels,
                                            (t_kernel_size, 1), # kernel_size
                                            (stride, 1), # stride
                                            ((t_kernel_size - 1) // 2, 0), # padding
                                            ),
                                  nn.BatchNorm2d(out_channels),
                                  nn.ReLU())
    def forward(self, x, A):
        x = self.t_gcn(self.s_gcn(x, A * self.M))
        return x

# STGCN网络模型
class STGCN(nn.Module):
    def __init__(self, A, num_classes, in_channels, t_kernel_size):
        super().__init__()
        self.A = A
        A_size = A.size()
        self.bn = nn.BatchNorm1d(in_channels * A_size[1])
        self.stgcn1 = STGCN_block(in_channels, 32, 1, t_kernel_size, A_size)
        self.stgcn2 = STGCN_block(32, 32, 1, t_kernel_size, A_size)
        self.stgcn3 = STGCN_block(32, 32, 1, t_kernel_size, A_size)
        self.stgcn4 = STGCN_block(32, 64, 2, t_kernel_size, A_size)
        self.stgcn5 = STGCN_block(64, 64, 1, t_kernel_size, A_size)
        self.stgcn6 = STGCN_block(64, 64, 1, t_kernel_size, A_size)
        self.fc = nn.Conv2d(64, num_classes, kernel_size=1)
    def forward(self, x):
        N, C, T, V = x.size()
        x = x.permute(0, 3, 1, 2).contiguous().view(N, V * C, T)
        x = self.bn(x)
        x = x.view(N, V, C, T).permute(0, 2, 3, 1).contiguous()
        x = self.stgcn1(x, self.A)
        x = self.stgcn2(x, self.A)
        x = self.stgcn3(x, self.A)
        x = self.stgcn4(x, self.A)
        x = self.stgcn5(x, self.A)
        x = self.stgcn6(x, self.A)
        x = F.avg_pool2d(x, x.size()[2:])
        x = x.view(N, -1, 1, 1)
        x = self.fc(x)
        x = x.view(x.size(0), -1)
        return x

def random_adjacency_matrix(num_nodes):
    upper_triangle = np.triu(np.random.randint(0, 2, size=(num_nodes, num_nodes)), k=1)
    diagonal = np.eye(num_nodes, dtype=int)
    adjacency_matrix = upper_triangle + upper_triangle.T + diagonal
    return adjacency_matrix

if __name__ == "__main__":
    num_nodes = 25
    adjacency_matrix = random_adjacency_matrix(num_nodes)
    adjacency_tensor_3d = torch.tensor(adjacency_matrix, dtype=torch.float32).unsqueeze(0)
    model = STGCN(adjacency_tensor_3d, 
                  num_classes=10,
                  in_channels=3,
                  t_kernel_size=9, 
                  )
    input = torch.randn(200, 3, 80, 25)
    output = model(input)
    print(output.shape)

输出结果为torch.Size([200, 10])，表示对应动作类别的概率。