[深度学习网络从入门到入土] 门控循环单元GRU

📢个人导航

知乎：https://www.zhihu.com/people/byzh_rc

CSDN：https://blog.csdn.net/qq_54636039

注：本文仅对所述内容做了框架性引导，具体细节可查询其余相关资料or源码

参考文章：各方资料

📖参考资料

Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translation.

🌱背景

经典 RNN 存在一个非常严重的问题：长序列训练时容易出现梯度消失 / 梯度爆炸

因此 LSTM 被提出，用 三个门 + memory cell 来解决长期依赖问题

但是 LSTM 结构较复杂：

• 参数多
• 计算量大
• 推理较慢

-> 一种更简单的结构：GRU（Gated Recurrent Unit）

核心思想: 用更少的门控结构，实现接近 LSTM 的效果

GRU 的主要改进：

• 只有 两个门
• 不再单独维护 cell state
• 参数更少, 训练更快

⚙️架构(公式)

1. 更新门(Update Gate) + 重置门(Reset Gate)

更新门控制：决定最终 hidden state 用多少新信息
$$z_t=\sigma (W_{xz}{x}_t+W_{hz}{h}_{t-1}+b_z)$$

• $z_t\approx 1$ → 保留旧记忆
• $z_t\approx 0$ → 使用新信息

重置门控制：决定生成新状态时要不要参考历史
$$r_t=\sigma (W_{xr}{x}_t+W_{hr}{h}_{t-1}+b_r)$$

• $r_t\approx 0$ → 忘掉过去信息
• $r_t\approx 1$ → 使用历史信息

2. 候选隐藏状态

候选状态：
$${\tilde{h}}_t=\tanh (W_{xh}{x}_t+W_{hh}(r_t\odot h_{t-1})+b_h)$$
注意：历史状态先经过 reset gate

门负责“控制比例”，所以用 sigmoid -> 限制在0~1
候选隐藏状态负责“产生内容”，所以用 tanh -> 有正有负, 没有偏移问题

3. 最终隐藏状态

最终输出由 update gate 决定：
$$h_t=(1-z_t)\odot h_{t-1}+z_t\odot {\tilde{h}}_t$$
可以理解为：旧状态 + 新状态的加权平均

4. GRU结构图

      h_{t-1}
        │
   ┌────┴────┐
   │         │
 update    reset
  gate      gate
   │         │
   │         ▼
   │      candidate
   │         │
   └────┬────┘
        ▼
       h_t

👍优点/创新点

1. 结构更简单

模型	门数量
RNN	0
GRU	2
LSTM	3

-> 参数更少, 计算更快 -> 收敛更快, 推理更快

2. 能学习长程依赖

通过 update gate：模型可以直接保留旧状态
$$h_t\approx h_{t-1}$$
这样可以缓解梯度消失问题

👎缺点

1. 表达能力略弱于 LSTM

GRU 没有独立的 cell state, 只有 $h_t$

-> 某些复杂任务中 LSTM 可能更强

2. 极长序列训练困难

这也是后来 Transformer 出现的重要原因

💻代码实现

import torch
import torch.nn as nn
import torch.nn.functional as F

from byzh.ai.Butils import b_get_params


class B_GRU_Paper(nn.Module):
    """
    单层 GRU（教学/论文公式对齐版）

    每个时间步:
        z_t = sigmoid(x_t W_xz + h_{t-1} W_hz + b_z)
        r_t = sigmoid(x_t W_xr + h_{t-1} W_hr + b_r)
        h~_t = tanh( x_t W_xh + (r_t * h_{t-1}) W_hh + b_h )

        h_t = (1 - z_t) * h_{t-1} + z_t * h~_t

    B = batch size
    T = sequence length
    D = input_size
    H = hidden_size
    C = output_size
    """

    def __init__(self, input_size: int, hidden_size: int,
                 output_size: int = None, batch_first: bool = True, bias: bool = True):
        super().__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.output_size = output_size
        self.batch_first = batch_first

        # 更新门: update gate
        self.x2z = nn.Linear(input_size, hidden_size, bias=bias)
        self.h2z = nn.Linear(hidden_size, hidden_size, bias=bias)

        # 重置门: reset gate
        self.x2r = nn.Linear(input_size, hidden_size, bias=bias)
        self.h2r = nn.Linear(hidden_size, hidden_size, bias=bias)

        # 候选隐藏状态: candidate hidden state
        self.x2h = nn.Linear(input_size, hidden_size, bias=bias)
        self.h2h = nn.Linear(hidden_size, hidden_size, bias=bias)

        # (可选)输出映射: optional output projection
        self.h2y = nn.Linear(hidden_size, output_size, bias=bias) if output_size is not None else None

        self.reset_parameters()

    def step(self, x_t, h_prev):
        """
        x_t:    (B, D)
        h_prev: (B, H)
        """
        z_t = torch.sigmoid(self.x2z(x_t) + self.h2z(h_prev))         # 更新门
        r_t = torch.sigmoid(self.x2r(x_t) + self.h2r(h_prev))         # 重置门
        h_hat = torch.tanh(self.x2h(x_t) + self.h2h(r_t * h_prev))    # 候选隐藏状态

        h_t = (1 - z_t) * h_prev + z_t * h_hat                        # 更新隐藏状态

        return h_t

    def forward(self, x, h0=None, return_sequences=True):
        """
        x:
            batch_first=True  -> (B, T, D)
            batch_first=False -> (T, B, D)

        h0:
            (B, H)

        return:
            hs: (B, T, H) or (T, B, H) or None
            hT: (B, H)
            ys: (B, T, C) or (T, B, C) or None
        """
        if not self.batch_first:
            x = x.transpose(0, 1)  # -> (B, T, D)

        B, T, D = x.shape

        h_t = x.new_zeros(B, self.hidden_size) if h0 is None else h0  # (B, H)

        hs = [] if return_sequences else None
        ys = [] if (return_sequences and self.h2y is not None) else None

        for t in range(T):
            x_t = x[:, t, :]
            h_t = self.step(x_t, h_t)

            if return_sequences:
                hs.append(h_t)
                if self.h2y is not None:
                    ys.append(self.h2y(h_t))

        hT = h_t

        if return_sequences:
            hs = torch.stack(hs, dim=1)  # (B, T, H)
            if ys is not None:
                ys = torch.stack(ys, dim=1)  # (B, T, C)

            if not self.batch_first:
                hs = hs.transpose(0, 1)
                if ys is not None:
                    ys = ys.transpose(0, 1)

        return hs, hT, ys

    def reset_parameters(self):
        modules = [
            self.x2z, self.h2z,
            self.x2r, self.h2r,
            self.x2h, self.h2h
        ]

        for m in modules:
            nn.init.xavier_uniform_(m.weight)
            if m.bias is not None:
                nn.init.zeros_(m.bias)

        if self.h2y is not None:
            nn.init.xavier_uniform_(self.h2y.weight)
            if self.h2y.bias is not None:
                nn.init.zeros_(self.h2y.bias)

class B_GRU_Paper_Layers(nn.Module):
    """
    多层 GRU（通过堆叠多个 B_GRU_Paper 实现）
    """

    def __init__(self, input_size, hidden_size, num_layers=2, output_size=None, batch_first=True):
        super().__init__()

        self.num_layers = num_layers
        self.batch_first = batch_first

        layers = []

        for i in range(num_layers):

            if i == 0:
                in_dim = input_size
            else:
                in_dim = hidden_size

            # 最后一层才接输出
            out_dim = output_size if i == num_layers - 1 else None

            layers.append(
                B_GRU_Paper(
                    input_size=in_dim,
                    hidden_size=hidden_size,
                    output_size=out_dim,
                    batch_first=batch_first
                )
            )

        self.layers = nn.ModuleList(layers)

    def forward(self, x):
        """
        return:
            hs: 最后一层所有时间步隐藏状态
            hT_list: 每一层最后时刻隐藏状态列表
            ys: 最后一层每个时间步输出（如果最后一层有 output_size）
        """
        hs = x
        hT_list = []
        ys = None

        for i, layer in enumerate(self.layers):
            hs, hT, ys = layer(hs, return_sequences=True)
            hT_list.append(hT)

        return hs, hT_list, ys


if __name__ == '__main__':
    # ===== 超参数 =====
    B = 50  # batch size
    T = 6  # sequence length
    D = 8  # input_size
    H = 16  # hidden_size
    C = 5  # output_size

    net = B_GRU_Paper(D, H, C, batch_first=True)
    a = torch.randn(B, T, D)
    hs, hT, ys = net(a)
    print(hT.shape)
    print(f"参数量: {b_get_params(net)}")  # 1_333

💻项目实例

库环境:

numpy==1.26.4
torch==2.2.2cu121
byzh-core==0.0.9.21
byzh-ai==0.0.9.61
byzh-extra==0.0.9.12
...

GRU训练MNIST数据集:

# copy all the codes from here to run

import torch
import torch.nn as nn
import torch.nn.functional as F

from byzh.ai.Btrainer import B_Classification_Trainer
from byzh.ai.Bdata import B_Download_MNIST, b_get_dataloader_from_tensor, b_stratified_indices
# from uploadToPypi_ai.byzh.ai.Bmodel.study_rnn import B_GRU_Paper
from byzh.ai.Bmodel.study_rnn import B_GRU_Paper
from byzh.ai.Butils import b_get_device

##### hyper params #####
epochs = 10
lr = 1e-3
batch_size = 128
device = b_get_device(use_idle_gpu=True)

##### data #####
downloader = B_Download_MNIST(save_dir='D:/study_model/datasets/MNIST')
data_dict = downloader.get_data()
X_train = data_dict['X_train_standard']
y_train = data_dict['y_train']
X_test = data_dict['X_test_standard']
y_test = data_dict['y_test']
num_classes = data_dict['num_classes']

train_dataloader, val_dataloader = b_get_dataloader_from_tensor(
    X_train, y_train, X_test, y_test,
    batch_size=batch_size
)

##### model #####
class MNIST_PixelGRU(nn.Module):
    """
    用 GRU 做 pixel-wise MNIST 分类。

    输入:
        x: (B, 1, 28, 28)

    处理流程:
        (B, 1, 28, 28)
        -> reshape
        -> (B, 28, 28)
        -> RNN
        -> 取最终隐藏状态 hT
        -> Linear(H, 10)
        -> logits: (B, 10)
    """

    def __init__(self, hidden_size=128, num_classes=10):
        super().__init__()

        self.gru = B_GRU_Paper(
            input_size=28,          # 每个时间步有28个像素值
            hidden_size=hidden_size,
            output_size=None,      # backbone 先不直接输出类别
            batch_first=True,
        )

        self.cls = nn.Linear(hidden_size, num_classes, bias=True)

        self.reset_parameters()

    def reset_parameters(self):
        nn.init.xavier_uniform_(self.cls.weight)
        if self.cls.bias is not None:
            nn.init.zeros_(self.cls.bias)

    def forward(self, x):
        """
        x: (B, 1, 28, 28)

        return:
            logits: (B, 10)
        """
        # (B, 1, 28, 28) -> (B, 28, 28)
        x = x.squeeze(1)

        # RNN 编码
        hs, hT, ys = self.gru(x, return_sequences=False)

        # 最终分类
        logits = self.cls(hT)  # (B, 10)
        return logits

model = MNIST_PixelGRU(num_classes=num_classes)

##### else #####
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
criterion = torch.nn.CrossEntropyLoss()

##### trainer #####
trainer = B_Classification_Trainer(
    model=model,
    optimizer=optimizer,
    criterion=criterion,
    train_loader=train_dataloader,
    val_loader=val_dataloader,
    device=device
)
trainer.set_writer1('./runs/gru/log.txt')

##### run #####
trainer.train_eval_s(epochs=epochs)

##### calculate #####
trainer.draw_loss_acc('./runs/gru/loss_acc.png', y_lim=False)
trainer.save_best_checkpoint('./runs/gru/best_checkpoint.pth')
trainer.calculate_model()