时空图神经网络3——T-GCN

啥也不会啥也白学

已于 2025-05-13 11:13:30 修改

阅读量1.2k

点赞数 23

CC 4.0 BY-SA版权

分类专栏：时空图神经网络文章标签：神经网络人工智能深度学习 gru 拓扑学

于 2025-05-04 14:28:17 首次发布

本文链接：https://blog.csdn.net/m0_51501260/article/details/147506345

时空图神经网络专栏收录该内容

6 篇文章

订阅专栏

系列文章目录

时序图神经网络1——GNN和GCN
时空图神经网络2——RNN和GRU

文章目录

系列文章目录
前言
一、方法介绍
总结

前言

本文是针对交通预测问题提出的一种考虑时间动态变化和城市道路拓扑结构的方法。我没有尝试复现，也没有使用，所以对本文提出方法的效果并不清楚，我不作推荐也不提出反对意见，只是在学习思路。
首先从时间的角度来说，交通流量具有周期性，和趋势性：
在这里插入图片描述
从空间的角度来说，上下游之间相互影响。上游的状态可以传递到下游，而下游的状态可以反馈到上游：

本文结合图卷积网络和门控循环单元提出了 T-GCN 模型。图卷积网络用于捕捉道路网络的拓扑结构，对空间依赖性进行建模。门控循环单元用于捕捉道路上交通数据的动态变化，对时间依赖性进行建模。T-GCN 模型还可应用于其他时空预测任务。
T-GCN 模型的预测结果在不同的预测时间跨度下均表现出稳定状态，这表明 T-GCN 模型不仅可以实现短期预测，还可用于长期交通流量预测任务。
本文使用深圳市罗湖区的出租车速度数据和洛杉矶高速公路（Los-loop）数据集对我们的方法进行评估。结果表明，与所有基准方法相比，我们的方法将预测误差降低了约 1.5%-57.8%，这证明了 T-GCN 模型在交通流量预测方面的优越性。

前面已经讲过GCN和GRU的内容，现在结合文章和源码详细拆解T-GCN。

参考文献：https://arxiv.org/abs/1811.05320
代码链接：https://github.com/lehaifeng/T-GCN/tree/master/T-GCN

一、方法介绍

1. 问题定义

本文中，交通流量预测的目标是基于道路上的历史交通信息，预测未来某一时间段内的交通信息。在我们的方法中，交通信息是一个通用概念，它可以是交通速度、交通流量和交通密度。不失一般性，在实验部分我们以交通速度作为交通信息的示例。
定义一，道路网络 $G$ ： 我们使用无向图 $G = (V, E)$ 来描述道路网络的拓扑结构，将每条道路视为一个节点。其中， $V$ 是道路节点的集合， $V={v_1 ,v_2,⋯,v_N}$ ， $N$ 是节点数量， $E$ 是边的集合。邻接矩阵 $A$ 用于表示道路之间的连接关系， $A\in R^{N×N}$ 。邻接矩阵的元素仅包含 0 和 1，元素为 0 表示道路之间没有连接，为 1 则表示存在连接。
定义二，特征矩阵 $X^{N×P}$ : 我们将道路网络上的交通信息视为网络中节点的属性特征，用 $\in R^{N×P}$ 表示。其中， $P$ 代表节点属性特征的数量（即历史时间序列的长度）， $X_t \in R^{N×i}$ 用于表示时刻i每条道路上的速度。同样，节点属性特征可以是任何交通信息，如交通速度、交通流量和交通密度。
因此，时空交通流量预测问题可以被看作是在道路网络拓扑结构 $G$ 和特征矩阵X的前提下，学习映射函数 $f$ ，然后计算接下来T个时刻的交通信息，如公式（1）所示：
$[X_{t + 1}, \cdots, X_{t + T}] = f(G; (X_{t - n}, \cdots, X_{t - 1}, X_t)) \quad(1)$
其中，n 是历史时间序列的长度，T 是需要预测的时间序列长度。
在这里插入图片描述
如上图所示，我们首先将历史的n个时间序列数据作为输入，利用图卷积网络来捕捉城市道路网络的拓扑结构，从而获取空间特征。其次，将得到的具有空间特征的时间序列数据输入到门控循环单元模型中，通过单元之间的信息传递来获取动态变化，进而捕捉时间特征。最后，我们通过全连接层得到预测结果。

2. 空间依赖性建模GCN

在这里插入图片描述
GCN的原理我们在之前讲过，在这里就不重复了，但是原文并没有使用pytorch里的GCN，我们现在对代码进行讲解。

class GCN(nn.Module):
    def __init__(self, adj, input_dim: int, output_dim: int, **kwargs):
        super(GCN, self).__init__()
        self.register_buffer(
            "laplacian", calculate_laplacian_with_self_loop(torch.FloatTensor(adj))
        )
        self._num_nodes = adj.shape[0]
        self._input_dim = input_dim  # seq_len for prediction
        self._output_dim = output_dim  # hidden_dim for prediction
        self.weights = nn.Parameter(
            torch.FloatTensor(self._input_dim, self._output_dim)
        )
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.xavier_uniform_(self.weights, gain=nn.init.calculate_gain("tanh"))

上面是GCN初始化部分，其中adj是邻接矩阵；input_dim是序列长度seq_len；output_dim是GCN的隐藏层尺度hidden_dim；register_buffer是 torch.nn.Module 类的一个方法，用于注册一个持久化的缓冲区(buffer)。缓冲区和参数(parameters)类似，但不被视为模型的可训练参数，即不会在反向传播过程中被更新。在这里，它将计算得到的图的拉普拉斯矩阵（通过calculate_laplacian_with_self_loop 函数）注册为名为 laplacian 的缓冲区。拉普拉斯矩阵听起来很复杂，实际上在讲GCN的时候已经讲过了，我们结合代码看一下：

def calculate_laplacian_with_self_loop(matrix):
    matrix = matrix + torch.eye(matrix.size(0))
    row_sum = matrix.sum(1)
    d_inv_sqrt = torch.pow(row_sum, -0.5).flatten()
    d_inv_sqrt[torch.isinf(d_inv_sqrt)] = 0.0
    d_mat_inv_sqrt = torch.diag(d_inv_sqrt)
    normalized_laplacian = (
        matrix.matmul(d_mat_inv_sqrt).transpose(0, 1).matmul(d_mat_inv_sqrt)
    )
    return normalized_laplacian

可以看出，该函数将传入的邻接矩阵（adj传入函数就是这里的matrix）的对角线上加1，目的是考虑节点本身（自循环）；再对 row_sum 中的每个元素进行 -0.5 次幂运算，然后将运算结果展平为一维张量，这一步用于计算度矩阵的逆平方根；下一行是把计算过程中可能出现的inf值替换成 0 （我觉得不会出现这种情况，因为对角线已经都加了 1 ，不存在度为 0 的节点，也就不存在inf）；torch.diag函数用于创建一个对角矩阵。它接受一个一维张量作为输入，将该张量的元素放在对角线上，其余位置填充为 0，从而得到一个方阵。这里将d_inv_sqrt转换为对角矩阵d_mat_inv_sqrt；matmul是 PyTorch 中用于矩阵乘法的函数，这里将原始矩阵 matrix与对角矩阵 d_mat_inv_sqrt 相乘；transpose(0, 1) 表示交换矩阵的行和列，在这里是进行转置。.matmul(d_mat_inv_sqrt)将转置后的矩阵再次与对角矩阵 d_mat_inv_sqrt 相乘，最终返回的就是 $\tilde{D}^{-\frac{1}{2}}\tilde{A}\tilde{D}^{-\frac{1}{2}}$ （这样看是不是就认识了）。
再回到初始化的代码，self._num_nodes是通过adj计算节点数；self.weights是生成的权重矩阵W；def reset_parameters(self)用于初始化模型的参数。nn.init.xavier_uniform_这是PyTorch中的一个参数初始化函数，它使用 Xavier 均匀分布来初始化张量。Xavier 初始化方法旨在保持输入和输出的方差在传播过程中尽可能一致，有助于缓解梯度消失和梯度爆炸问题。self.weights是要初始化的张量，即前面定义的权重矩阵。gain=nn.init.calculate_gain("tanh")，gain 是一个缩放因子，用于调整初始化分布的范围。nn.init.calculate_gain("tanh") 会根据激活函数 tanh 计算合适的缩放因子，确保在使用 tanh 激活函数时，初始化的参数能够更好地适应网络的训练。

    def forward(self, inputs):
        # (batch_size, seq_len, num_nodes)
        batch_size = inputs.shape[0]
        # (num_nodes, batch_size, seq_len)
        inputs = inputs.transpose(0, 2).transpose(1, 2)
        # (num_nodes, batch_size * seq_len)
        inputs = inputs.reshape((self._num_nodes, batch_size * self._input_dim))
        # AX (num_nodes, batch_size * seq_len)
        ax = self.laplacian @ inputs
        # (num_nodes, batch_size, seq_len)
        ax = ax.reshape((self._num_nodes, batch_size, self._input_dim))
        # (num_nodes * batch_size, seq_len)
        ax = ax.reshape((self._num_nodes * batch_size, self._input_dim))
        # act(AXW) (num_nodes * batch_size, output_dim)
        outputs = torch.tanh(ax @ self.weights)
        # (num_nodes, batch_size, output_dim)
        outputs = outputs.reshape((self._num_nodes, batch_size, self._output_dim))
        # (batch_size, num_nodes, output_dim)
        outputs = outputs.transpose(0, 1)
        return outputs

看懂了__init__，forward部分的原理就没什么好说的了，主要是看数据的尺寸是如何变化的，这对我们自己写代码也很有帮助。输入input的尺寸是[batch_size, seq_len, num_nodes]，把num_nodes换成num_features的话就跟我们之前讲过的RNN是一样的（batch_first=True），毕竟我们这是处理时序数据，所以比纯GCN多了一维。然后通过transpose把尺寸变成[num_nodes, batch_size, seq_len]，这样做是为了方便后续与图的拉普拉斯矩阵进行矩阵乘法，拉普拉斯矩阵的形状是 [num_nodes, num_nodes]，调整后可以使矩阵乘法在节点维度上进行。然后再变成[num_nodes, batch_size*seq_len]，这是为了进行矩阵乘法，需要将输入张量转换为二维张量，以便与拉普拉斯矩阵进行乘法运算。算完了以后再展成[num_nodes, batch_size, seq_len]，再次将 [num_nodes] 维和 [batch_size] 维合并为一维，将张量形状变为 [num_nodes * batch_size, seq_len]，这是为了与权重矩阵 self.weights（形状为 [self._input_dim, self._output_dim]）进行矩阵乘法（在__init__中已经提过了，input_dim和seq_len是一样的）。最后再展开，再换位，最终输出的尺寸是[batch_size, num_nodes, output_dim]。

3. 时间依赖性模型GRU

在这里插入图片描述
GRU的原理我们在之前讲过，在这里就不重复了，但是原文并没有使用pytorch里的GRU，我们现在对代码进行讲解：
GRU分为三个部分GRULinear，GRUCell和GRU。

3.1 GRU

我们从外向内讲，最外部的就是GRU。

class GRU(nn.Module):
    def __init__(self, input_dim: int, hidden_dim: int, **kwargs):
        super(GRU, self).__init__()
        self._input_dim = input_dim  # num_nodes for prediction
        self._hidden_dim = hidden_dim
        self.gru_cell = GRUCell(self._input_dim, self._hidden_dim)

    def forward(self, inputs):
        batch_size, seq_len, num_nodes = inputs.shape
        assert self._input_dim == num_nodes
        outputs = list()
        hidden_state = torch.zeros(batch_size, num_nodes * self._hidden_dim).type_as(
            inputs
        )
        for i in range(seq_len):
            output, hidden_state = self.gru_cell(inputs[:, i, :], hidden_state)
            output = output.reshape((batch_size, num_nodes, self._hidden_dim))
            outputs.append(output)
        last_output = outputs[-1]
        return last_output

从forward第一行可以看出，输入的尺寸是[batch_size, seq_len, num_nodes]，没有什么意外，并且在第二行检查是否self._input_dim == num_nodes，也就是说，在初始化__init__的时候输入的input_dim实际上就是num_node。隐藏状态的尺寸是[batch_size, num_nodes*self._hidden_dim]，这是为什么呢？我们后面再看。接下来对seq_len进行循环，也就是说，我们让batch中的第i条数据以及i-1的隐藏层经过GRUCell，得到的输出储存起来，得到的隐藏层用在第i+1条数据的输入。得到的结果还要经过一次维度的变换，我们在讲GCN的时候讲过，输出跟隐藏层的尺寸是一样的，所以要把输出再展开成三维，这样我们才能知道每个batch中的每个node里的具体计算结果是什么。
另外，我们看输入到GRUCell的两个输入尺寸分别是[batch_size, num_nodes]（三维变二维）和[batch_size, num_nodes*self._hidden_dim]（本来就是二维），这样可以回答刚才的问题了。我们知道x的第三维是num_features，也就是特征数，我们在深度学习中一般也都是对特征进行线性变换，也就是num_features变成hidden_dim。在GRU中类似，只不过这个位置不是对特征进行变换，而是对节点中的性质进行变换，变换后的大小应该可以被num_nodes整除，所以使用num_nodes*self._hidden_dim作为隐藏层的第二维。

3.2 GRUCell

上面GRU的代码中可以看出，输入和隐藏层是经过GRUCell得到最终结果的，我们现在就看这一部分呢起到了什么作用。

class GRUCell(nn.Module):
    def __init__(self, input_dim: int, hidden_dim: int):
        super(GRUCell, self).__init__()
        self._input_dim = input_dim
        self._hidden_dim = hidden_dim
        self.linear1 = GRULinear(self._hidden_dim, self._hidden_dim * 2, bias=1.0)
        self.linear2 = GRULinear(self._hidden_dim, self._hidden_dim)

    def forward(self, inputs, hidden_state):
        # [r, u] = sigmoid([x, h]W + b)
        # [r, u] (batch_size, num_nodes * (2 * num_gru_units))
        concatenation = torch.sigmoid(self.linear1(inputs, hidden_state))
        # r (batch_size, num_nodes * num_gru_units)
        # u (batch_size, num_nodes * num_gru_units)
        r, u = torch.chunk(concatenation, chunks=2, dim=1)
        # c = tanh([x, (r * h)]W + b)
        # c (batch_size, num_nodes * num_gru_units)
        c = torch.tanh(self.linear2(inputs, r * hidden_state))
        # h := u * h + (1 - u) * c
        # h (batch_size, num_nodes * num_gru_units)
        new_hidden_state = u * hidden_state + (1 - u) * c
        return new_hidden_state, new_hidden_state

其实代码中的注释都写得很清楚了，看过上一篇blog的话，应该熟悉重置门r，更新门u和候选隐藏状态c，但是我们还是来捋一下。forward中第一行注释说明拼接起来的r和u的计算公式，是拼接起来的输入和隐藏状态进行线性变换再经过sigmoid，concatenation这一行能看出其代码实现。第二行注释说明拼接起来的r和u的尺寸，这个尺寸我们从初始化的linear1中也能看出来，但是他的名称有点变化，从_hidden_dim变成了num_gru_units，这我不懂为什么要改名。torch.chunk(concatenation, chunks=2, dim=1)将 concatenation 张量在第 1 维上分割成两个部分，每个部分的形状为 [batch_size, num_nodes * hidden_dim]，也就是把拼接起来的r 和 u 拆开。从这里要拆开我们推测GRULinear中应该是将输入和隐藏状态拼接在一起用同一个矩阵进行线性变换，这跟分别线性变化再加起来没有区别，很多GRU的讲解也都是这么讲的。再经过GRULinear得到候选隐藏状态，最后得到隐藏状态。
在讲GRU的时候我就说，看到公式以后应该就能够把代码写出来，所以最后我也懒得仔。这段代码实际上就是公式的直译，所见即所得，作为初学者（指我自己）应该敢于动手写。

3.3 GRULinear

实现了一个自定义的线性变换模块，用于在 GRU 单元中对输入和隐藏状态进行拼接和线性变换，在GRUCell中用到两次。

class GRULinear(nn.Module):
    def __init__(self, num_gru_units: int, output_dim: int, bias: float = 0.0):
        super(GRULinear, self).__init__()
        self._num_gru_units = num_gru_units
        self._output_dim = output_dim
        self._bias_init_value = bias
        self.weights = nn.Parameter(
            torch.FloatTensor(self._num_gru_units + 1, self._output_dim)
        )
        self.biases = nn.Parameter(torch.FloatTensor(self._output_dim))
        self.reset_parameters()

    def forward(self, inputs, hidden_state):
        batch_size, num_nodes = inputs.shape
        # inputs (batch_size, num_nodes, 1)
        inputs = inputs.reshape((batch_size, num_nodes, 1))
        # hidden_state (batch_size, num_nodes, num_gru_units)
        hidden_state = hidden_state.reshape(
            (batch_size, num_nodes, self._num_gru_units)
        )
        # [inputs, hidden_state] "[x, h]" (batch_size, num_nodes, num_gru_units + 1)
        concatenation = torch.cat((inputs, hidden_state), dim=2)
        # [x, h] (batch_size * num_nodes, gru_units + 1)
        concatenation = concatenation.reshape((-1, self._num_gru_units + 1))
        # [x, h]W + b (batch_size * num_nodes, output_dim)
        outputs = concatenation @ self.weights + self.biases
        # [x, h]W + b (batch_size, num_nodes, output_dim)
        outputs = outputs.reshape((batch_size, num_nodes, self._output_dim))
        # [x, h]W + b (batch_size, num_nodes * output_dim)
        outputs = outputs.reshape((batch_size, num_nodes * self._output_dim))
        return outputs

如果直接看这一部分代码可能还不明白啥是num_gru_unit，现在我们知道了，其实就是_hidden_dim。这边的代码尺寸变化也写得很清楚。在GRUCell中我们知道第一个入参接收的永远是输入，所以尺寸一直都是[batch_size, num_nodes]；第二个入参接收的是隐藏状态或者隐藏状态与其相同尺寸矩阵的 $\odot$ 积，所以尺寸一直都是[batch_size, num_nodes, num_gru_units]。让输入加一维到三维，把隐藏状态展从二维开成三维，再把它们的第二维（0, 1, 2，也就是最后一维）拼接在一起，尺寸就变成了[batch_size, num_nodes, num_gru_units+1]；然后把三维的前两维合在一起变成二维方便矩阵运算，尺寸变成[batch_size*num_nodes, num_gru_units+1]；经过一次线性变换再改两次尺寸输出（前前后后的各种变换都是为了这次线性变换，为了这醋才包的饺子）。

4. 时间图卷积网络

在这里插入图片描述
左侧展示了时空交通预测的过程，右侧展示了 T-GCN 单元的具体结构。在结构中， $h_{t-1}$ 表示 $t - 1$ 时刻的输出，GC 代表图卷积过程， $u_t$ 和 $r_t$ 分别是t时刻的更新门和重置门， $h_t$ 表示 $t$ 时刻的输出。这些元素共同构成了 T-GCN 单元的结构，展示了数据在模型中的流动和处理方式。
我们可以看出，时空图卷积模型的一个单元是GCN和GRU的结合。

先通过GCN；
再把卷积结果和上一时刻的GRU计算的隐藏层输入本时刻的GRU；
得到新的隐藏层用于下一时刻的计算，还得到当前时刻的输出。

TGCN公式如下：
在这里插入图片描述
TGCN的代码也是分为三个部分：TGCN、TGCNCell和TGCNGraphConvolution，看过前面的部分再看这些应该会简单很多。

4.1 TGCN

class TGCN(nn.Module):
    def __init__(self, adj, hidden_dim: int, **kwargs):
        super(TGCN, self).__init__()
        self._input_dim = adj.shape[0]
        self._hidden_dim = hidden_dim
        self.register_buffer("adj", torch.FloatTensor(adj))
        self.tgcn_cell = TGCNCell(self.adj, self._input_dim, self._hidden_dim)

    def forward(self, inputs):
        batch_size, seq_len, num_nodes = inputs.shape
        assert self._input_dim == num_nodes
        hidden_state = torch.zeros(batch_size, num_nodes * self._hidden_dim).type_as(
            inputs
        )
        output = None
        for i in range(seq_len):
            output, hidden_state = self.tgcn_cell(inputs[:, i, :], hidden_state)
            output = output.reshape((batch_size, num_nodes, self._hidden_dim))
        return output

这一部分非常简单，初始化的时候输入邻接矩阵adj和隐藏层维度。模型输入的尺寸是[batch_size, seq_len, num_nodes]，前面GCN和GRU的输入也一直都是这个。后面的代码也跟GRU的代码基本一样。但是一样就引来一个问题，因为我们这里已经对seq_len循环了，后面到了GRU还怎么继续循环？我们接着看。

4.2 GRUCell

class TGCNCell(nn.Module):
    def __init__(self, adj, input_dim: int, hidden_dim: int):
        super(TGCNCell, self).__init__()
        self._input_dim = input_dim
        self._hidden_dim = hidden_dim
        self.register_buffer("adj", torch.FloatTensor(adj))
        self.graph_conv1 = TGCNGraphConvolution(
            self.adj, self._hidden_dim, self._hidden_dim * 2, bias=1.0
        )
        self.graph_conv2 = TGCNGraphConvolution(
            self.adj, self._hidden_dim, self._hidden_dim
        )

    def forward(self, inputs, hidden_state):
        # [r, u] = sigmoid(A[x, h]W + b)
        # [r, u] (batch_size, num_nodes * (2 * num_gru_units))
        concatenation = torch.sigmoid(self.graph_conv1(inputs, hidden_state))
        # r (batch_size, num_nodes, num_gru_units)
        # u (batch_size, num_nodes, num_gru_units)
        r, u = torch.chunk(concatenation, chunks=2, dim=1)
        # c = tanh(A[x, (r * h)W + b])
        # c (batch_size, num_nodes * num_gru_units)
        c = torch.tanh(self.graph_conv2(inputs, r * hidden_state))
        # h := u * h + (1 - u) * c
        # h (batch_size, num_nodes * num_gru_units)
        new_hidden_state = u * hidden_state + (1.0 - u) * c
        return new_hidden_state, new_hidden_state

我发现这段跟GRU不还是一样吗，就是把linear改成了卷积，这让我意识到一些不对劲，立马回到代码最上面看导入的包：

import argparse
import torch
import torch.nn as nn
from utils.graph_conv import calculate_laplacian_with_self_loop

笑死，根本没用到刚才讲的GRU和GCN，我才意识到这是作者用于对比效果的代码，这三个之间是相互独立的，我可以直接讲TGCN的，对不起大家（私密马喽.JPG）！但是看完GRU和GCN以后再看就TGCN就没有障碍了，多学了知识，这是好事啊！所以我就不修改前面的内容了，一起浪费时间吧！

4.3 TGCNGraphConvolution

class TGCNGraphConvolution(nn.Module):
    def __init__(self, adj, num_gru_units: int, output_dim: int, bias: float = 0.0):
        super(TGCNGraphConvolution, self).__init__()
        self._num_gru_units = num_gru_units
        self._output_dim = output_dim
        self._bias_init_value = bias
        self.register_buffer(
            "laplacian", calculate_laplacian_with_self_loop(torch.FloatTensor(adj))
        )
        self.weights = nn.Parameter(
            torch.FloatTensor(self._num_gru_units + 1, self._output_dim)
        )
        self.biases = nn.Parameter(torch.FloatTensor(self._output_dim))
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.xavier_uniform_(self.weights)
        nn.init.constant_(self.biases, self._bias_init_value)

    def forward(self, inputs, hidden_state):
        batch_size, num_nodes = inputs.shape
        # inputs (batch_size, num_nodes) -> (batch_size, num_nodes, 1)
        inputs = inputs.reshape((batch_size, num_nodes, 1))
        # hidden_state (batch_size, num_nodes, num_gru_units)
        hidden_state = hidden_state.reshape(
            (batch_size, num_nodes, self._num_gru_units)
        )
        # [x, h] (batch_size, num_nodes, num_gru_units + 1)
        concatenation = torch.cat((inputs, hidden_state), dim=2)
        # [x, h] (num_nodes, num_gru_units + 1, batch_size)
        concatenation = concatenation.transpose(0, 1).transpose(1, 2)
        # [x, h] (num_nodes, (num_gru_units + 1) * batch_size)
        concatenation = concatenation.reshape(
            (num_nodes, (self._num_gru_units + 1) * batch_size)
        )
        # A[x, h] (num_nodes, (num_gru_units + 1) * batch_size)
        a_times_concat = self.laplacian @ concatenation
        # A[x, h] (num_nodes, num_gru_units + 1, batch_size)
        a_times_concat = a_times_concat.reshape(
            (num_nodes, self._num_gru_units + 1, batch_size)
        )
        # A[x, h] (batch_size, num_nodes, num_gru_units + 1)
        a_times_concat = a_times_concat.transpose(0, 2).transpose(1, 2)
        # A[x, h] (batch_size * num_nodes, num_gru_units + 1)
        a_times_concat = a_times_concat.reshape(
            (batch_size * num_nodes, self._num_gru_units + 1)
        )
        # A[x, h]W + b (batch_size * num_nodes, output_dim)
        outputs = a_times_concat @ self.weights + self.biases
        # A[x, h]W + b (batch_size, num_nodes, output_dim)
        outputs = outputs.reshape((batch_size, num_nodes, self._output_dim))
        # A[x, h]W + b (batch_size, num_nodes * output_dim)
        outputs = outputs.reshape((batch_size, num_nodes * self._output_dim))
        return outputs