STGCN(时空图卷积网络)详解

1️⃣ STGCN介绍

前面已经介绍过了图卷积（GCN）。这篇论文《Spatial Temporal Graph Convolutional Networks for Skeleton-Based Action Recognition》将GCN扩展到时空图模型上，用于实现动作识别。下图展示了STGCN的输入，即一系列骨架图。其中每个节点对应于人体的一个关节，有两种类型的边，①符合节点自然连通性的空间边（图1中淡蓝色线条）②跨越连续时间步长连接相同节点的时间边（淡绿色线条）

---

2️⃣ 网络结构

STGCN简单的网络结构如下所示，由三部分构成：

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

在这里，我们不关注数据部分，只对网络结构进行解析，因此我们来看ST-GCN块的详细结构：

• 步骤一：引入一个可学习的权重矩阵，与邻接矩阵大小一致，记作Learnable edge importance weight。让它与邻接矩阵$A$按位相乘，得到加权后的邻接矩阵。其目的是给重要的边较大的权重，给非重要的边较小的权重。
• 步骤二：将加权后的邻接矩阵与输入数据送到GCN中进行运算
• 步骤三：利用TCN网络，实现时间维度信息的聚合

---

3️⃣ 代码

# 只包含网络结构，“网络的输入（关节坐标）”和“邻接矩阵”都是随机生成的

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


# ---------------------------------------------------------------------------------
# 空域图卷积
# ---------------------------------------------------------------------------------
class SpatialGraphConvolution(nn.Module):
    """
    Args:
    in_channels: 输入通道数，表示每个节点的特征维度（例如 3 对应 x, y, z 坐标）。
    out_channels: 输出通道数，表示经过卷积后每个节点的特征维度。
    s_kernel_size: 空间卷积核的大小，等于邻接矩阵的数量，表示多重图卷积的支持
    """
    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)
        #对邻接矩阵进行GC，相加特征
        x = torch.einsum('nkctv,kvw->nctw', (x, A))
        return x.