GRU层的计算过程是怎样的？PyTorch nn.GRU运算维度错误排查

GRU计算错误的核心原因 #

你定义的GRU隐状态输出维度为96，GRU包含重置门、更新门、候选隐状态三个独立的门控结构，每个门都对应一组独立的权重和偏置参数。PyTorch为了提升计算效率，会把三个门的参数沿第一个维度拼接存储，3*96=288就是你看到所有参数第一维都是288的原因。你计算时没有拆分这三部分参数，直接用全量288维的结果作为单门的输出，才会得到错误的维度。

正确计算步骤 #

1. 拆分拼接的参数 #

把所有参数沿第一维拆为3等份，每份对应一个门的参数，维度如下：

• weight_ih_l0拆为W_ir、W_iz、W_in，单份维度均为(96,96)
• weight_hh_l0拆为W_hr、W_hz、W_hn，单份维度均为(96,96)
• bias_ih_l0拆为b_ir、b_iz、b_in，单份维度均为(96,)
• bias_hh_l0拆为b_hr、b_hz、b_hn，单份维度均为(96,)

2. 单时间步计算逻辑 #

你取出的单时间步输入x_t维度为(B=8, Hin=96)，上一步隐状态h_{t-1}维度为(B=8, H_out=96)，按以下公式计算即可得到正确维度的结果：

1. 计算重置门r_t，维度为(8,96)：

r_t = sigmoid( x_t @ W_ir.T + b_ir + h_{t-1} @ W_hr.T + b_hr )

2. 计算更新门z_t，维度为(8,96)：

z_t = sigmoid( x_t @ W_iz.T + b_iz + h_{t-1} @ W_hz.T + b_hz )

3. 计算候选隐状态n_t，维度为(8,96)，其中*为哈达玛积：

n_t = tanh( x_t @ W_in.T + b_in + r_t * (h_{t-1} @ W_hn.T + b_hn) )

4. 计算当前步隐状态h_t，维度为(8,96)，和要求的隐状态维度完全一致：

h_t = (1 - z_t) * n_t + z_t * h_{t-1}

3. 3维输入的计算说明 #

你提到的权重为2维、输入为3维的问题，PyTorch内部会自动做广播计算：输入维度(L=1000, B=8, 96)和转置后的weight_ih（96,288）做矩阵乘法时，会自动把前两个维度（序列长度、批次大小）作为批量维度，仅对最后一维做矩阵运算，输出维度为(1000, 8, 288)，之后再拆分288为三个96维分别计算三个门即可，不需要手动逐时间步计算。

内容的提问来源于stack exchange，提问作者void