pytorch RNNCell

pytorch RNNCell的一次前馈本质上就是一个线性层。如图所示：

可以看到，每一次把输入送入RNNCell的过程，就是把输入部分和隐藏状态部分分别经过简单全连接层，并接上激活函数的过程。

再附上RNNCell简易demo便于理解：

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

# 配置
input_size = 4       # 每步输入维度
hidden_size = 16     # 隐藏层维度
seq_len = 6          # 序列长度
batch_size = 8       # 批大小
num_classes = 2      # 分类数量
epochs = 30

# 定义模型
class RNNCellClassifier(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        super().__init__()
        self.rnn_cell = nn.RNNCell(input_size, hidden_size)
        self.fc = nn.Linear(hidden_size, num_classes)

    def forward(self, x):
        # x: [seq_len, batch_size, input_size]
        h_t = torch.zeros(x.size(1), hidden_size, device=x.device)
        for t in range(x.size(0)):
            h_t = self.rnn_cell(x[t], h_t)  # [batch_size,hidden_size]
        out = self.fc(h_t)  # 最后一个隐藏状态用于分类
        return out

# 生成训练数据
def generate_batch(batch_size, seq_len, input_size):
    x = torch.randn(seq_len, batch_size, input_size)
    # 以最后一个时间步第一个特征是否 > 0 来构造 label
    last_step = x[-1]
    labels = (last_step[:, 0] > 0).long()
    return x, labels

# 模型初始化
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = RNNCellClassifier(input_size, hidden_size, num_classes).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)

# 训练
for epoch in range(epochs):
    model.train()
    x_batch, y_batch = generate_batch(batch_size, seq_len, input_size)
    x_batch, y_batch = x_batch.to(device), y_batch.to(device)

    logits = model(x_batch)
    loss = criterion(logits, y_batch)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if (epoch + 1) % 5 == 0 or epoch == 0:
        pred = logits.argmax(dim=1)
        acc = (pred == y_batch).float().mean().item()
        print(f"[Epoch {epoch+1}] Loss: {loss.item():.4f}, Acc: {acc:.2f}")

# 测试一组
model.eval()
with torch.no_grad():
    x_test, y_test = generate_batch(1, seq_len, input_size)
    x_test, y_test = x_test.to(device), y_test.to(device)
    pred = model(x_test).argmax(dim=1)
    print("\n Test sample:")
    print("Target label:", y_test.item())
    print("Predicted   :", pred.item())