PyTorch 中使用截断反向传播 (BPTT) 训练 RNN 单元

import torch
import torch.nn as nn
import torch.optim as optim

# 定义 RNN 单元
class RNNCell(nn.Module):
    def __init__(self, input_size, hidden_size):
        super(RNNCell, self).__init__()
        self.hidden_size = hidden_size
        self.i2h = nn.Linear(input_size + hidden_size, hidden_size)
        self.i2o = nn.Linear(input_size + hidden_size, 1) # 输出大小设置为1，简化示例

    def forward(self, input, hidden):
        combined = torch.cat((input, hidden), 1)
        hidden = torch.tanh(self.i2h(combined))
        output = torch.sigmoid(self.i2o(combined)) # 使用sigmoid激活函数
        return output, hidden

# 超参数
input_size = 10
hidden_size = 20
sequence_length = 100
batch_size = 32
bptt_length = 20
learning_rate = 0.01
num_epochs = 10

# 初始化 RNN 单元和优化器
rnn_cell = RNNCell(input_size, hidden_size)
optimizer = optim.Adam(rnn_cell.parameters(), lr=learning_rate)
criterion = nn.BCELoss() # 使用二元交叉熵损失函数

# 训练数据 (示例)
data = torch.randn(sequence_length, batch_size, input_size)
targets = torch.randint(0, 2, (sequence_length, batch_size, 1)).float() # 二分类目标

# 训练循环
for epoch in range(num_epochs):
    # 初始化隐藏状态
    hidden = torch.zeros(batch_size, hidden_size)

    # 遍历序列，以 bptt_length 为步长
    for i in range(0, sequence_length - bptt_length, bptt_length):
        # 截取一个块
        inputs = data[i:i+bptt_length]
        target_batch = targets[i:i+bptt_length]

        # 清零梯度
        optimizer.zero_grad()

        # 前向传播和计算损失
        loss = 0
        for j in range(bptt_length):
            output, hidden = rnn_cell(inputs[j], hidden)
            loss += criterion(output, target_batch[j])

        # 反向传播
        loss.backward()

        # 梯度裁剪 (可选，但推荐)
        torch.nn.utils.clip_grad_norm_(rnn_cell.parameters(), 5) # 梯度裁剪阈值

        # 更新参数
        optimizer.step()

        # 分离隐藏状态，防止梯度传播到之前的块
        hidden = hidden.detach()

    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')