基于pytorch的RNN实现字符级姓氏文本分类的示例代码

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import numpy as np

class SurnameDataset(Dataset):
    def __init__(self, data):
        self.data = data
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        return self.data[idx]
        
# 假设我们的数据格式为 (surname, country)，例如 ('Smith', 'USA')
# 这里假设数据已经预处理成对应的数值表示
# 例如将字符映射为数字，国家名称映射为数字等
 
# 数据预处理函数
def preprocess_data(data):
    processed_data = []
    for surname, country in data:
        # 将姓氏转换为字符索引列表
        surname_indices = [char_to_index[char] for char in surname]
        # 将国家转换为对应的数字
        country_index = country_to_index[country]
        processed_data.append((surname_indices, country_index))
    return processed_data

class SurnameRNN(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(SurnameRNN, self).__init__()
        self.hidden_size = hidden_size
        self.embedding = nn.Embedding(input_size, hidden_size)
        self.lstm = nn.LSTM(hidden_size, hidden_size)
        self.fc = nn.Linear(hidden_size, output_size)
 
    def forward(self, input, hidden):
        embedded = self.embedding(input).view(1, 1, -1)
        output, hidden = self.lstm(embedded, hidden)
        output = self.fc(output.view(1, -1))
        return output, hidden
 
    def init_hidden(self):
        return (torch.zeros(1, 1, self.hidden_size), torch.zeros(1, 1, self.hidden_size))

input_size = len(char_to_index)  # 姓氏中字符的数量
hidden_size = 128
output_size = len(country_to_index)  # 国家的数量
learning_rate = 0.001
num_epochs = 10
 
model = SurnameRNN(input_size, hidden_size, output_size)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=learning_rate)
 
# 假设我们有一个经过预处理的数据集 surname_data
# 数据格式为 (surname_indices, country_index)
 
# 将数据划分为训练集和测试集
train_data = surname_data[:800]
test_data = surname_data[800:]
 
# 开始训练
for epoch in range(num_epochs):
    total_loss = 0
    for surname_indices, country_index in train_data:
        model.zero_grad()
        hidden = model.init_hidden()
        surname_tensor = torch.tensor(surname_indices, dtype=torch.long)
        country_tensor = torch.tensor([country_index], dtype=torch.long)
 
        for i in range(len(surname_indices)):
            output, hidden = model(surname_tensor[i], hidden)
        
        loss = criterion(output, country_tensor)
        total_loss += loss.item()
        loss.backward()
        optimizer.step()
    
    print('Epoch [{}/{}], Loss: {:.4f}'.format(epoch+1, num_epochs, total_loss / len(train_data)))

# 导入必要的库
import torch
from torch.utils.data import DataLoader
 
# 定义测试函数
def test_model(model, test_data):
    model.eval()  # 设置模型为评估模式
    correct = 0
    total = 0
    with torch.no_grad():
        for surname_indices, country_index in test_data:
            surname_tensor = torch.tensor(surname_indices, dtype=torch.long)
            country_tensor = torch.tensor([country_index], dtype=torch.long)
            
            hidden = model.init_hidden()
            
            for i in range(len(surname_indices)):
                output, hidden = model(surname_tensor[i], hidden)
            
            _, predicted = torch.max(output.data, 1)
            
            total += 1
            if predicted == country_tensor:
                correct += 1
    
    accuracy = correct / total
    print('Accuracy on test data: {:.2%}'.format(accuracy))
 
# 加载测试数据集
test_dataset = SurnameDataset(test_data)
test_loader = DataLoader(test_dataset, batch_size=1, shuffle=True)
 
# 加载已经训练好的模型
model_path = "path_to_your_trained_model.pt"
model = torch.load(model_path)
 
# 测试模型
test_model(model, test_loader)