01-循环神经网络基础¶

import torch
import torch.nn as nn

class ManualRNNCell(nn.Module):  # 继承nn.Module定义神经网络层
    """手动实现 RNN Cell"""
    def __init__(self, input_size, hidden_size):  # __init__构造方法，创建对象时自动调用
        super().__init__()  # super()调用父类方法
        self.hidden_size = hidden_size

        # 参数
        self.W_xh = nn.Parameter(torch.randn(input_size, hidden_size) * 0.01)
        self.W_hh = nn.Parameter(torch.randn(hidden_size, hidden_size) * 0.01)
        self.b_h = nn.Parameter(torch.zeros(hidden_size))

    def forward(self, x_t, h_prev):
        """
        x_t: (batch, input_size) 当前时间步输入
        h_prev: (batch, hidden_size) 上一个隐藏状态
        """
        h_t = torch.tanh(x_t @ self.W_xh + h_prev @ self.W_hh + self.b_h)
        return h_t

class ManualRNN(nn.Module):
    """手动实现完整 RNN"""
    def __init__(self, input_size, hidden_size, num_classes):
        super().__init__()
        self.hidden_size = hidden_size
        self.rnn_cell = ManualRNNCell(input_size, hidden_size)
        self.fc = nn.Linear(hidden_size, num_classes)

    def forward(self, x, h_0=None):
        """
        x: (batch, seq_len, input_size)
        """
        batch_size, seq_len, _ = x.shape

        if h_0 is None:
            h_0 = torch.zeros(batch_size, self.hidden_size, device=x.device)

        h_t = h_0
        outputs = []
        for t in range(seq_len):
            h_t = self.rnn_cell(x[:, t, :], h_t)
            outputs.append(h_t)

        # outputs: list of (batch, hidden_size)
        outputs = torch.stack(outputs, dim=1)  # (batch, seq_len, hidden_size)

        # 取最后一个时间步的输出进行分类
        out = self.fc(h_t)
        return out, outputs

# 测试
model = ManualRNN(input_size=10, hidden_size=64, num_classes=5)
x = torch.randn(32, 20, 10)  # batch=32, seq_len=20, input=10
out, all_hidden = model(x)
print(f"输出: {out.shape}, 所有隐藏状态: {all_hidden.shape}")

def truncated_bptt_training(model, data_stream, seq_length=35, optimizer=None, criterion=None):
    """截断 BPTT 训练"""
    h = None
    for i in range(0, len(data_stream) - 1, seq_length):
        # 获取片段
        x = data_stream[i:i+seq_length]
        y = data_stream[i+1:i+seq_length+1]

        # 截断梯度链，但保留隐藏状态值
        if h is not None:
            h = h.detach()  # detach()从计算图分离，不参与梯度计算

        output, h = model(x, h)
        loss = criterion(output.view(-1, output.size(-1)), y.view(-1))  # view重塑张量形状（要求内存连续）

        optimizer.zero_grad()  # 清零梯度，防止梯度累积
        loss.backward()  # 反向传播计算梯度
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0)
        optimizer.step()  # 根据梯度更新模型参数

class ManualLSTMCell(nn.Module):
    """手动实现 LSTM Cell"""
    def __init__(self, input_size, hidden_size):
        super().__init__()
        self.hidden_size = hidden_size

        # 将四个门的参数合并为一个大矩阵以提高效率
        self.linear_ih = nn.Linear(input_size, 4 * hidden_size)
        self.linear_hh = nn.Linear(hidden_size, 4 * hidden_size)

    def forward(self, x_t, state):
        """
        x_t: (batch, input_size)
        state: (h_prev, c_prev) 各 (batch, hidden_size)
        """
        h_prev, c_prev = state

        # 一次计算四个门
        gates = self.linear_ih(x_t) + self.linear_hh(h_prev)

        # 分割
        i_gate, f_gate, g_gate, o_gate = gates.chunk(4, dim=1)

        # 激活
        i_t = torch.sigmoid(i_gate)      # 输入门
        f_t = torch.sigmoid(f_gate)      # 遗忘门
        g_t = torch.tanh(g_gate)         # 候选细胞
        o_t = torch.sigmoid(o_gate)      # 输出门

        # 更新细胞状态
        c_t = f_t * c_prev + i_t * g_t

        # 计算隐藏状态
        h_t = o_t * torch.tanh(c_t)

        return h_t, c_t

class ManualLSTM(nn.Module):
    """手动实现完整 LSTM"""
    def __init__(self, input_size, hidden_size, num_layers=1):
        super().__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers

        self.cells = nn.ModuleList([
            ManualLSTMCell(
                input_size if i == 0 else hidden_size,
                hidden_size
            )
            for i in range(num_layers)
        ])

    def forward(self, x, initial_state=None):
        """
        x: (batch, seq_len, input_size)
        """
        batch_size, seq_len, _ = x.shape

        if initial_state is None:
            h = [torch.zeros(batch_size, self.hidden_size, device=x.device) for _ in range(self.num_layers)]  # 列表推导式，简洁创建列表
            c = [torch.zeros(batch_size, self.hidden_size, device=x.device) for _ in range(self.num_layers)]
        else:
            h, c = initial_state
            h = [h[i] for i in range(self.num_layers)]
            c = [c[i] for i in range(self.num_layers)]

        outputs = []
        for t in range(seq_len):
            layer_input = x[:, t, :]
            for layer_idx, cell in enumerate(self.cells):  # enumerate同时获取索引和元素
                h[layer_idx], c[layer_idx] = cell(layer_input, (h[layer_idx], c[layer_idx]))
                layer_input = h[layer_idx]
            outputs.append(h[-1])  # [-1]负索引取最后一个元素

        outputs = torch.stack(outputs, dim=1)
        return outputs, (torch.stack(h), torch.stack(c))

# 使用 PyTorch 内置 LSTM
lstm = nn.LSTM(
    input_size=100,    # 输入特征维度
    hidden_size=256,   # 隐藏状态维度
    num_layers=2,      # 层数
    batch_first=True,  # 输入 shape: (batch, seq, feature)
    dropout=0.5,       # 层间 dropout
    bidirectional=False
)

x = torch.randn(32, 50, 100)
output, (h_n, c_n) = lstm(x)
print(f"输出: {output.shape}")  # (32, 50, 256)
print(f"最终隐藏: {h_n.shape}") # (2, 32, 256) — num_layers × batch × hidden
print(f"最终细胞: {c_n.shape}") # (2, 32, 256)

# PyTorch GRU
gru = nn.GRU(
    input_size=100,
    hidden_size=256,
    num_layers=2,
    batch_first=True,
    dropout=0.5
)

x = torch.randn(32, 50, 100)
output, h_n = gru(x)
print(f"GRU 输出: {output.shape}")    # (32, 50, 256)
print(f"GRU 隐藏: {h_n.shape}")       # (2, 32, 256)

# 双向 LSTM
bilstm = nn.LSTM(
    input_size=100,
    hidden_size=256,
    num_layers=2,
    batch_first=True,
    bidirectional=True  # 双向
)

x = torch.randn(32, 50, 100)
output, (h_n, c_n) = bilstm(x)
print(f"BiLSTM 输出: {output.shape}")  # (32, 50, 512) — 512 = 256 * 2
print(f"BiLSTM h_n: {h_n.shape}")     # (4, 32, 256) — 4 = 2layers * 2directions

class StackedLSTM(nn.Module):
    """带层间dropout的多层LSTM"""
    def __init__(self, input_size, hidden_size, num_layers, dropout=0.5):
        super().__init__()
        self.layers = nn.ModuleList()
        self.dropouts = nn.ModuleList()

        for i in range(num_layers):
            self.layers.append(nn.LSTM(
                input_size if i == 0 else hidden_size,
                hidden_size,
                batch_first=True
            ))
            if i < num_layers - 1:
                self.dropouts.append(nn.Dropout(dropout))

    def forward(self, x):
        for i, layer in enumerate(self.layers):
            x, _ = layer(x)
            if i < len(self.dropouts):
                x = self.dropouts[i](x)
        return x

class TextClassifier(nn.Module):
    """基于LSTM的文本分类模型"""
    def __init__(self, vocab_size, embed_dim, hidden_dim, num_classes,
                 num_layers=2, bidirectional=True, dropout=0.5, pad_idx=0):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=pad_idx)
        self.lstm = nn.LSTM(
            embed_dim, hidden_dim, num_layers,
            batch_first=True, bidirectional=bidirectional, dropout=dropout
        )
        self.dropout = nn.Dropout(dropout)
        direction = 2 if bidirectional else 1
        self.fc = nn.Linear(hidden_dim * direction, num_classes)

    def forward(self, x, lengths=None):
        # x: (batch, seq_len)
        embedded = self.dropout(self.embedding(x))

        if lengths is not None:
            packed = nn.utils.rnn.pack_padded_sequence(
                embedded, lengths.cpu(), batch_first=True, enforce_sorted=False
            )
            output, (h_n, c_n) = self.lstm(packed)
            output, _ = nn.utils.rnn.pad_packed_sequence(output, batch_first=True)
        else:
            output, (h_n, c_n) = self.lstm(embedded)

        # 取最后时间步 (前向最后 + 后向最后)
        if self.lstm.bidirectional:
            hidden = torch.cat([h_n[-2], h_n[-1]], dim=1)
        else:
            hidden = h_n[-1]

        return self.fc(self.dropout(hidden))

# 使用
model = TextClassifier(
    vocab_size=10000, embed_dim=300, hidden_dim=256,
    num_classes=2, num_layers=2, bidirectional=True
)
x = torch.randint(0, 10000, (32, 100))
lengths = torch.randint(10, 100, (32,))
output = model(x, lengths)
print(f"分类输出: {output.shape}")  # (32, 2)