导入所需库

import torch import torch.nn as nn import numpy as np

定义数据集路径和一些超参数

data_path = '/path/to/pinyin/data' batch_size = 32 num_epochs = 10 input_size = 128 # 输入向量大小 hidden_size = 256 # 隐状态向量大小 num_layers = 2 # 神经网络层数 learning_rate = 0.001

定义数据预处理函数

def process_data(file_path): ''' 读取拼音数据文件，将其转化为one-hot编码的形式 ''' with open(file_path, 'r', encoding='utf-8') as f: data = f.read().strip().split('\n') char_to_idx = dict() # 字符到索引的映射 for line in data: pinyin, word = line.split('\t') pinyin = pinyin.split(' ') for c in pinyin: if c not in char_to_idx: char_to_idx[c] = len(char_to_idx) idx_to_char = {i: c for c, i in char_to_idx.items()} # 索引到字符的映射 # 将拼音转化为one-hot编码 X = [] Y = [] for line in data: pinyin, word = line.split('\t') pinyin = pinyin.split(' ') x = [char_to_idx[c] for c in pinyin] # 将拼音转化为索引 y = [char_to_idx[c] for c in pinyin[1:]] + [char_to_idx['']] # 标签Y为X向右移一位，最后一位为结束标记 x = np.eye(len(char_to_idx))[x] # one-hot编码 y = np.eye(len(char_to_idx))[y] # one-hot编码 X.append(x) Y.append(y) return X, Y, char_to_idx, idx_to_char

加载数据

X, Y, char_to_idx, idx_to_char = process_data(data_path)

定义数据生成器函数

def data_generator(X, Y, batch_size, shuffle=True): num_batches = len(X) // batch_size # 计算批次数 indices = list(range(len(X))) if shuffle: np.random.shuffle(indices) for i in range(num_batches): batch_indices = indices[i * batch_size : (i + 1) * batch_size] X_batch = [X[idx] for idx in batch_indices] Y_batch = [Y[idx] for idx in batch_indices] X_batch = torch.tensor(X_batch, dtype=torch.float32) Y_batch = torch.tensor(Y_batch, dtype=torch.float32) yield X_batch, Y_batch

定义循环神经网络模型

class RNN(nn.Module): def init(self, input_size, hidden_size, num_layers, output_size): super(RNN, self).init() self.hidden_size = hidden_size self.num_layers = num_layers self.rnn = nn.GRU(input_size, hidden_size, num_layers, batch_first=True) self.fc = nn.Linear(hidden_size, output_size)

def forward(self, x, h0=None):
    if h0 is None:
        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size)
    out, h = self.rnn(x, h0)
    out = self.fc(out)
    return out, h

定义模型、损失函数和优化器

model = RNN(input_size, hidden_size, num_layers, len(char_to_idx)) criterion = nn.CrossEntropyLoss() optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

训练模型

total_step = len(X) // batch_size for epoch in range(num_epochs): for i, (x, y) in enumerate(data_generator(X, Y, batch_size)): # 前向传播 outputs, _ = model(x) loss = 0 for j in range(outputs.size(1)): loss += criterion(outputs[:, j, :], torch.argmax(y[:, j, :], dim=-1)) # 反向传播和优化 optimizer.zero_grad() loss.backward() nn.utils.clip_grad_norm_(model.parameters(), max_norm=1) # 梯度截断 optimizer.step() # 打印训练信息 if (i+1) % 100 == 0: print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}' .format(epoch+1, num_epochs, i+1, total_step, loss.item()))

测试模型

prefix = 'ni hao' # 给定一个前缀 prefix = prefix.split(' ') x = [char_to_idx[c] for c in prefix] x = np.eye(len(char_to_idx))[x] x = torch.tensor(x, dtype=torch.float32).unsqueeze(0) with torch.no_grad(): # 单步预测 output, _ = model(x) output = output[0][-1] predict = torch.argmax(output).item() print(prefix + [idx_to_char[predict]]) # K步预测 for i in range(10): output, _ = model(x) output = output[0][-1] predict = torch.argmax(output).item() print(prefix + [idx_to_char[predict]]) x = np.concatenate([x, np.eye(len(char_to_idx))[predict].reshape(1, 1, -1)], axis=1) x = torch.tensor(x, dtype=torch.float32)