在 PyTorch 中实现 GRU（门控循环单元）非常简单，它的用法和 LSTM 很相似，但结构更简洁，参数更少。下面用 “文本生成” 的例子（根据前几个词续写句子），一步步讲解 GRU 的实现方法和关键注意事项。

一、先搞懂 PyTorch 中 GRU 的 “基本用法”

PyTorch 提供了nn.GRU类，封装了 GRU 的核心逻辑，我们只需要关注几个关键参数：

import torch.nn as nn

# 定义一个GRU层
gru = nn.GRU(
    input_size=100,    # 每个输入的特征数（比如每个词用100维向量表示）
    hidden_size=200,   # 隐藏层大小（记忆单元的维度，自己设定）
    num_layers=2,      # GRU的层数（2层就是深层GRU）
    batch_first=True   # 输入数据格式是否为 [batch_size, 序列长度, 特征数]
)

• 输入格式：假设我们有 32 个句子（batch_size=32），每句 10 个词（序列长度 = 10），每个词用 100 维向量表示，输入形状就是 (32, 10, 100)。
• 输出格式：GRU 会返回两个结果：
  1. 所有时间步的输出：(32, 10, 200)（每一步的隐藏状态）
  2. 最后一个时间步的隐藏状态：(2, 32, 200)（因为有 2 层，第一维是层数）

这里要注意：GRU 比 LSTM 简单，没有 “细胞状态”，只返回隐藏状态（因为 GRU 把 LSTM 的 “细胞状态” 和 “隐藏状态” 合并了）。

二、完整实现：用 GRU 做文本生成

我们以 “根据前几个词续写句子” 为例，展示 GRU 的完整实现流程。

1. 准备数据：把文字转换成数字序列

import torch

# 训练文本（简单的句子集合）
texts = ["我 爱 吃 苹果", "今天 天气 很 好", "小明 在 看书", "春天 花儿 开 了"]
# 构建词汇表（所有出现的词）
words = [word for text in texts for word in text.split()]
word_to_id = {word: i for i, word in enumerate(list(set(words)))}
id_to_word = {i: word for word, i in word_to_id.items()}
vocab_size = len(word_to_id)

# 把文本转换成输入-目标对（比如“我 爱 吃”作为输入，“苹果”作为目标）
def create_dataset(texts, seq_len=3):
    inputs, targets = [], []
    for text in texts:
        ids = [word_to_id[word] for word in text.split()]
        # 生成多个输入-目标对
        for i in range(len(ids) - seq_len):
            inputs.append(ids[i:i+seq_len])  # 前3个词作为输入
            targets.append(ids[i+seq_len])   # 第4个词作为目标
    return torch.tensor(inputs, dtype=torch.long), torch.tensor(targets, dtype=torch.long)

# 生成训练数据
inputs, targets = create_dataset(texts, seq_len=3)  # inputs形状: (n, 3), targets形状: (n,)

2. 搭建 GRU 模型：词嵌入 + GRU + 输出层

GRU 模型的结构和 LSTM 类似，但输出部分更简单（因为没有细胞状态）：

PyTorch实现GRU文本生成模型

import torch
import torch.nn as nn

class GRUGenerator(nn.Module):
    def __init__(self, vocab_size, embedding_dim, hidden_dim, num_layers):
        super(GRUGenerator, self).__init__()
        # 1. 词嵌入层：把词ID转换成向量
        self.embedding = nn.Embedding(
            num_embeddings=vocab_size,  # 词汇表大小
            embedding_dim=embedding_dim  # 词向量维度
        )
        
        # 2. GRU层
        self.gru = nn.GRU(
            input_size=embedding_dim,   # 输入特征数（和词向量维度一致）
            hidden_size=hidden_dim,     # 隐藏层大小
            num_layers=num_layers,      # GRU层数
            batch_first=True,           # 输入格式为 [batch_size, seq_len, embedding_dim]
            dropout=0.2                 # 防止过拟合（可选）
        )
        
        # 3. 输出层：把GRU的输出转换成词汇表中每个词的概率
        self.fc = nn.Linear(hidden_dim, vocab_size)
    
    def forward(self, x, hidden=None):
        # 步骤1：词嵌入，形状从 [batch_size, seq_len] → [batch_size, seq_len, embedding_dim]
        x_embed = self.embedding(x)  # 例如：(8,3) → (8,3,100)
        
        # 步骤2：通过GRU
        # out: 所有时间步的输出，形状 [batch_size, seq_len, hidden_dim]
        # h_n: 最后一个时间步的隐藏状态，形状 [num_layers, batch_size, hidden_dim]
        # 如果没提供初始hidden，GRU会自动用全0初始化
        out, h_n = self.gru(x_embed, hidden)
        
        # 步骤3：取最后一个时间步的输出，转换成词概率
        # 最后一个时间步的输出形状：[batch_size, hidden_dim]
        last_out = out[:, -1, :]
        logits = self.fc(last_out)  # 形状: [batch_size, vocab_size]
        
        return logits, h_n

3. 训练模型：让 GRU 学会续写句子

# 模型参数
embedding_dim = 100  # 词向量维度
hidden_dim = 128     # GRU隐藏层大小
num_layers = 2       # GRU层数

# 初始化模型、损失函数和优化器
model = GRUGenerator(vocab_size, embedding_dim, hidden_dim, num_layers)
criterion = nn.CrossEntropyLoss()  # 多分类损失（因为要预测词汇表中的词）
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# 训练500轮
for epoch in range(500):
    # 前向传播
    logits, _ = model(inputs)
    loss = criterion(logits, targets)
    
    # 反向传播 + 优化
    optimizer.zero_grad()  # 清空梯度
    loss.backward()        # 计算梯度
    optimizer.step()       # 更新参数
    
    # 每100轮打印一次损失
    if (epoch+1) % 100 == 0:
        print(f'Epoch [{epoch+1}/500], Loss: {loss.item():.4f}')

4. 测试模型：用 GRU 续写句子

def generate_text(start_text, max_len=5):
    # 把起始文本转换成ID
    start_ids = [word_to_id[word] for word in start_text.split()]
    input_seq = torch.tensor([start_ids], dtype=torch.long)  # 形状: (1, seq_len)
    
    generated = start_text.split()
    hidden = None  # 初始隐藏状态
    
    for _ in range(max_len):
        # 预测下一个词
        logits, hidden = model(input_seq, hidden)
        # 取概率最大的词作为预测结果
        next_id = torch.argmax(logits, dim=1).item()
        next_word = id_to_word[next_id]
        generated.append(next_word)
        
        # 更新输入序列（去掉第一个词，加入新预测的词）
        input_seq = torch.cat([input_seq[:, 1:], torch.tensor([[next_id]])], dim=1)
    
    return ' '.join(generated)

# 测试：用“我 爱 吃”开头续写
print(generate_text("我 爱 吃"))  # 可能输出："我 爱 吃 苹果 了"（取决于训练效果）

三、关键注意事项：这些细节要牢记

1. GRU 和 LSTM 的主要区别在输出

   • LSTM 返回 (out, (h_n, c_n))（包含细胞状态 c_n）；
   • GRU 只返回 (out, h_n)（没有细胞状态，因为合并了）。
     所以在提取最终状态时，GRU 直接用 h_n[-1, :, :] 即可（取最后一层的隐藏状态）。

2. 初始化隐藏状态的方式
   默认情况下，GRU 会自动用全 0 初始化隐藏状态，但也可以手动指定（比如用特定分布初始化）：

   # 手动初始化隐藏状态（形状：[num_layers, batch_size, hidden_dim]）
   h0 = torch.randn(num_layers, batch_size, hidden_dim)
   out, hn = gru(x_embed, h0)  # 传入GRU
   

3. 处理变长序列时的技巧
   和 LSTM 一样，遇到长度不同的句子（用 0 填充）时，需要用pack_padded_sequence忽略填充的 0，否则 GRU 会把 0 当成有效输入：

   from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
   
   # 假设lengths是每个句子的真实长度
   packed = pack_padded_sequence(x_embed, lengths, batch_first=True, enforce_sorted=False)
   out_packed, hn = gru(packed)  # GRU会跳过填充的0
   out, _ = pad_packed_sequence(out_packed, batch_first=True)  # 恢复原形状
   

4. GRU 的优势：训练更快，适合资源有限的场景
   GRU 的参数比 LSTM 少约 1/3（没有细胞状态和输出门），所以：

   • 训练时占用显存更少（适合 GPU 内存小的设备）；
   • 收敛速度更快（相同数据下，GRU 可能比 LSTM 少用 30% 的训练时间）。
     实际项目中，如果 LSTM 训练太慢，可以试试 GRU，多数场景下效果差异很小。

5. 避免过拟合的小技巧

   • 增加dropout：在 GRU 参数中设置dropout=0.2（层之间的 dropout）；
   • 减少隐藏层大小：比如把hidden_dim从 256 降到 128；
   • 早停（Early Stopping）：当验证集损失不再下降时，提前停止训练。

总结

PyTorch 实现 GRU 的核心步骤是：
“数据预处理（文字→数字→向量）→ 搭建模型（嵌入层 + GRU 层 + 输出层）→ 训练优化→ 生成 / 预测”

GRU 的用法和 LSTM 高度相似，但结构更简单（少了细胞状态），计算更快，适合对速度和显存要求较高的场景。记住 “输出只含隐藏状态”“处理变长序列用 pack_padded_sequence” 这两个关键点，就能轻松用好 GRU 处理长序列任务了。