英译法案例 - 测试基于 GRU 的编码器

import torch
import torch.nn as nn


class EncoderGRU(nn.Module):
    """基于 GRU 的编码器"""
    
    def __init__(self, vocab_size, hidden_size):
        super(EncoderGRU, self).__init__()
        self.hidden_size = hidden_size
        
        # 1. 词嵌入层
        self.embedding = nn.Embedding(vocab_size, hidden_size)
        
        # 2. GRU 层
        self.gru = nn.GRU(hidden_size, hidden_size)
    
    def forward(self, input, hidden):
        # 1. 词嵌入
        embedded = self.embedding(input).view(1, 1, -1)
        
        # 2. GRU
        output, hidden = self.gru(embedded, hidden)
        
        return output, hidden
    
    def initHidden(self):
        """初始化隐藏状态"""
        return torch.zeros(1, 1, self.hidden_size)


def test_encoder_gru():
    """测试基于 GRU 的编码器"""
    print("=== 英译法案例 - 测试基于 GRU 的编码器 ===")
    
    # 1. 获取数据加载器对象
    # my_data_loader = get_data_loader()
    
    # 2. 初始化编码器参数
    vocab_size = 28036  # 英文词汇表大小
    hidden_size = 256   # 隐藏层维度
    
    # 3. 创建编码器模型对象
    my_encoder_gru = EncoderGRU(vocab_size, hidden_size)
    
    # 4. 将模型移动到指定设备（CPU/GPU）
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    my_encoder_gru = my_encoder_gru.to(device)
    
    print(f"设备: {device}")
    print(f"词汇表大小: {vocab_size}")
    print(f"隐藏层维度: {hidden_size}")
    
    # 5. 遍历数据加载器进行测试
    # for i, (x, y) in enumerate(my_data_loader):
    
    # 模拟测试
    input_tensor = torch.tensor([[4243, 139, 7907, 4323, 9088, 1]])
    print(f"\n输入张量形状: {input_tensor.shape}")
    print(f"输入张量内容: {input_tensor}")
    
    # 6. 前向传播
    hidden = my_encoder_gru.initHidden().to(device)
    input_tensor = input_tensor.to(device)
    
    for i in range(input_tensor.size(1)):
        input_word = torch.tensor([[input_tensor[0][i].item()]]).to(device)
        output, hidden = my_encoder_gru(input_word, hidden)
        print(f"单词 {i+1} - output 形状: {output.shape}, hidden 形状: {hidden.shape}")
    
    print(f"\n最终 hidden 形状: {hidden.shape}")


# 测试
if __name__ == "__main__":
    test_encoder_gru()