import torch
import torch
.nn as nn
import torch
.optim as optim
import numpy as np
import collections
import os
import time
import matplotlib
.pyplot as plt
# 设置超参数
start_token = 'G' # 诗歌起始标记
end_token = 'E' # 诗歌结束标记
batch_size = 64 # 训练批量大小
embedding_dim = 128 # 词向量维度
hidden_dim = 256 # LSTM隐藏层维度
learning_rate =
0.001 # 学习率
num_epochs = 5
0 # 训练轮数
# 设备配置（自动选择GPU或CPU）
device = torch
.device
("cuda" if torch
.cuda
.is_available
() else "cpu"
)
print
(f"Using device: {device}"
)
# 词嵌入层
class WordEmbedding
(nn
.Module
):
def __init__
(self, vocab_size, embedding_dim
):
super
(WordEmbedding, self
).__init__
()
self
.embedding = nn
.Embedding
(vocab_size, embedding_dim
)
nn
.init
.uniform_
(self
.embedding
.weight, -1
.0, 1
.0) # 均匀初始化
def forward
(self, x
):
return self
.embedding
(x
)
# RNN模型
class RNN_Model
(nn
.Module
):
def __init__
(self, vocab_size, embedding_dim, hidden_dim, num_layers=2, dropout=
0.3
):
super
(RNN_Model, self
).__init__
()
self
.embedding = WordEmbedding
(vocab_size, embedding_dim
)
self
.lstm = nn
.LSTM
(
input_size=embedding_dim,
hidden_size=hidden_dim,
num_layers=num_layers,
batch_first=True,
dropout=dropout if num_layers > 1 else
0
)
self
.fc = nn
.Linear
(hidden_dim, vocab_size
)
self
.dropout = nn
.Dropout
(dropout
)
self
.init_weights
()
def init_weights
(self
):
# 全连接层初始化
nn
.init
.xavier_uniform_
(self
.fc
.weight
)
nn
.init
.zeros_
(self
.fc
.bias
)
# LSTM权重初始化
for name, param in self
.lstm
.named_parameters
():
if 'weight' in name:
nn
.init
.orthogonal_
(param
)
elif 'bias' in name:
nn
.init
.zeros_
(param
)
# 设置遗忘门偏置为1（有助于缓解梯度消失）
n = param
.size
(0)
param
.data[n // 4:n // 2]
.fill_
(1
.0)
def forward
(self, x, hidden=None
):
# x:
(batch_size, seq_len
)
batch_size = x
.size
(0)
# 嵌入层
embeds = self
.embedding
(x
) #
(batch_size, seq_len, embedding_dim
)
# LSTM层
if hidden is None:
h
0 = torch
.zeros
(self
.lstm
.num_layers, batch_size, self
.lstm
.hidden_size
).to
(device
)
c
0 = torch
.zeros
(self
.lstm
.num_layers, batch_size, self
.lstm
.hidden_size
).to
(device
)
hidden =
(h
0, c
0)
lstm_out, hidden = self
.lstm
(embeds, hidden
) #
(batch_size, seq_len, hidden_dim
)
# 应用dropout，防止过拟合
lstm_out = self
.dropout
(lstm_out
)
# 全连接层
output = self
.fc
(lstm_out
) #
(batch_size, seq_len, vocab_size
)
# 重新排列维度用于损失计算
output = output
.permute
(0, 2, 1
) #
(batch_size, vocab_size, seq_len
)
return output, hidden
# 数据处理函数
def process_poems2
(file_name
):
poems = []
with open
(file_name, "r", encoding='utf-8'
) as f:
for line in f
.readlines
():
try:
line = line
.strip
()
if line:
content = line
.replace(' '' ', ''
).replace('，', ''
).replace('。', ''
)
if '_' in content or '
(' in content or '（' in content or '《' in content or '[' in content or \
start_token in content or end_token in content:
continue
if len
(content
) < 5 or len
(content
) > 8
0:
continue
content = start_token + content + end_token
poems
.append
(content
)
except ValueError as e:
pass
# 按诗的字数排序
poems = sorted
(poems,
key=lambda line: len
(line
))
# 统计每个字出现次数
all_
words = []
for poem in poems:
all_
words += [word for word in poem]
counter = collections
.Counter
(all_
words) # 统计词和词频。
count_pairs = sorted
(counter
.items
(),
key=lambda x: -x[1]
) # 排序
words, _ = zip
(*count_pairs
)
words =
words[:len
(words)] +
(' ',
)
word_int_
map = dict
(zip
(words, range
(len
(words))))
poems_vector = [list
(map(word_int_
map.get, poem
)) for poem in poems]
return poems_vector, word_int_
map,
words
# 批量生成函数
def create_batches
(poems_idx, batch_size, max_seq_len=5
0):
# 按长度排序,有助于减少填充
sorted_poems = sorted
(poems_idx,
key=len
)
batches = []
num_batches = len
(sorted_poems
) // batch_size
for i in range
(num_batches
):
start_idx = i * batch_size
end_idx =
(i + 1
) * batch_size
batch = sorted_poems[start_idx:end_idx]
# 找到本批次最大长度
max_len = min
(max
(len
(poem
) for poem in batch
), max_seq_len
)
# 填充序列
padded_batch = []
for poem in batch:
if len
(poem
) > max_len:
padded_poem = poem[:max_len] # 截断
else:
padded_poem = poem + [
0] *
(max_len - len
(poem
)) # 填充，
0是<PAD>
padded_batch
.append
(padded_poem
)
# 输入和输出序列
inputs = [poem[:-1] for poem in padded_batch] # 输入序列（去除最后一个字符）
targets = [poem[1:] for poem in padded_batch] # 输出序列（移除第一个字符）
batches
.append
((inputs, targets
))
return batches
# 训练函数
def train_model
():
# 处理数据
poems_idx, char_to_idx, idx_to_char = process_poems2
('
./tangshi
.txt'
)
vocab_size = len
(char_to_idx
)
print
(f"Vocabulary size: {vocab_size}"
)
print
(f"Number of poems: {len
(poems_idx
)}"
)
# 创建批次
batches = create_batches
(poems_idx, batch_size
)
print
(f"Number of batches: {len
(batches
)}"
)
# 初始化模型
model = RNN_Model
(
vocab_size=vocab_size,
embedding_dim=embedding_dim,
hidden_dim=hidden_dim,
num_layers=2
).to
(device
)
# 模型加载
pretrained_model_path = '
./models/best_poetry_model'
best_model_path = '
./models/best_tangshi_model
.pth'
# 确保模型目录存在
os
.makedirs
('
./models', exist_ok=True
)
# 检查预训练模型是否能加载
model_loaded = False
if os
.path
.exists
(pretrained_model_path
):
try:
# 尝试加载预训练模型
model
.load_state_dict
(torch
.load
(pretrained_model_path
))
print
(f"Loaded pretrained model from: {pretrained_model_path}"
)
model_loaded = True
except Exception as e:
print
(f"Error loading pretrained model: {e}"
)
if not model_loaded:
print
("No valid pretrained model found
. Starting training from scratch
."
)
# 损失函数和优化器
criterion = nn
.CrossEntropyLoss
(ignore_index=
0) # 忽略填充索引
optimizer = optim
.Adam
(model
.parameters
(), lr=learning_rate
)
scheduler = optim
.lr_scheduler
.ReduceLROnPlateau
(
optimizer, mode='min', factor=
0.5, patience=3, verbose=True
)
# 训练记录
train_losses = []
best_loss = float
('inf'
)
print
("Starting training
..."
)
start_time = time
.time
()
for epoch in range
(num_epochs
):
model
.train
()
epoch_loss =
0.0
for batch_idx,
(inputs, targets
) in enumerate
(batches
):
# 转换为张量
inputs_tensor = torch
.tensor
(inputs, dtype=torch
.long
).to
(device
)
targets_tensor = torch
.tensor
(targets, dtype=torch
.long
).to
(device
)
# 前向传播
optimizer
.zero_grad
()
output, _ = model
(inputs_tensor
)
# 计算损失
loss = criterion
(output, targets_tensor
)
# 反向传播
loss
.backward
()
# 梯度裁剪（防止梯度爆炸）
nn
.utils
.clip_grad_norm_
(model
.parameters
(), 1
.0)
# 更新参数
optimizer
.step
()
# 记录损失
epoch_loss += loss
.item
()
# 打印进度
if batch_idx % 2
0 ==
0:
avg_loss = epoch_loss /
(batch_idx + 1
)
elapsed = time
.time
() - start_time
print
(f"Epoch [{epoch + 1}/{num_epochs}], Batch [{batch_idx}/{len
(batches
)}], "
f"Loss: {avg_loss:
.4f}, Time: {elapsed:
.2f}s"
)
# 计算本轮平均损失
epoch_loss /= len
(batches
)
train_losses
.append
(epoch_loss
)
scheduler
.step
(epoch_loss
)
# 打印摘要
print
(f"Epoch [{epoch + 1}/{num_epochs}] completed, Avg Loss: {epoch_loss:
.4f}"
)
# 保存最佳模型
if epoch_loss < best_loss:
best_loss = epoch_loss
torch
.save
({
'epoch': epoch,
'model_state_dict': model
.state_dict
(),
'optimizer_state_dict': optimizer
.state_dict
(),
'loss': epoch_loss,
'char_to_idx': char_to_idx,
'idx_to_char': idx_to_char
}, best_model_path
)
print
(f"Saved best model with loss: {best_loss:
.4f} at {best_model_path}"
)
# 绘制损失曲线
plt
.figure
(figsize=
(1
0, 5
))
plt
.plot
(train_losses, label='Training Loss'
)
plt
.title
('Training Loss Over Epochs'
)
plt
.xlabel
('Epochs'
)
plt
.ylabel
('Loss'
)
plt
.legend
()
plt
.savefig
('
./models/tangshi_training_loss
.png'
)
plt
.close
()
print
("Training completed!"
)
# 生成诗歌函数
def generate_poem
(model, idx_to_char, char_to_idx, start_word, max_length=4
0):
model
.eval
()
poem = []
# 初始化输入
input_seq = torch
.tensor
([[char_to_idx[start_token], char_to_idx[start_word]]],
dtype=torch
.long
).to
(device
)
poem
.extend
([start_token, start_word]
)
# 初始化隐藏状态
hidden = None
with torch
.no_grad
(): # 禁用梯度计算
for _ in range
(max_length
):
# 向前传播，预测下一个字符
output, hidden = model
(input_seq, hidden
)
# 获取最后一个字符的预测
last_output = output[:, :, -1] #
(batch_size, vocab_size
)
# 应用温度采样/多项式采样（增加多样性）
probabilities = torch
.softmax
(last_output, dim=1
).squeeze
()
next_idx = torch
.multinomial
(probabilities, 1
).item
()
# 检查是否结束
if next_idx == char_to_idx[end_token]:
break
# 添加到诗歌中
next_char = idx_to_char
.get
(next_idx, '<UNK>'
)
poem
.append
(next_char
)
# 更新输入序列
input_seq = torch
.tensor
([[next_idx]], dtype=torch
.long
).to
(device
)
return ''
.join
(poem
)
# 打印格式化的诗歌
def pretty_print_poem
(poem
):
# 移除标记并分割句子
clean_poem = poem
.replace(start_token, ''
).replace(end_token, ''
)
sentences = clean_poem
.split
('。'
)
# 打印非空句子
for s in sentences:
if s
.strip
():
print
(s
.strip
() + '。'
)
# 保存诗歌到文件
def save_poems_to_file
(poems, filename="generated_poems
.txt"
):
with open
(filename, 'w', encoding='utf-8'
) as f:
f
.write
("=== 生成古诗 ===\n\n"
)
for i,
(word, poem
) in enumerate
(poems
):
f
.write
(f"诗歌 {i + 1}
(起始词: {word}
)\n"
)
clean_poem = poem
.replace(start_token, ''
).replace(end_token, ''
)
# 按句号分行
for char in clean_poem:
f
.write
(char
)
if char in ['，', '。', '！', '？']:
f
.write
('\n'
)
f
.write
("\n---------------------\n\n"
)
print
(f"诗歌已保存到: {filename}"
)
# 主函数
if __name__ == '__main__':
# 确保模型目录存在
os
.makedirs
('
./models', exist_ok=True
)
# 检查是否需要训练
best_model_path = '
./models/best_tangshi_model
.pth'
if not os
.path
.exists
(best_model_path
):
print
("训练模型未找到，开始训练
..."
)
train_model
()
else:
print
("找到已有的训练模型，跳过训练"
)
# 加载最佳模型并生成诗歌
print
("\nGenerating poems with best model
..."
)
try:
# 尝试加载模型
checkpoint = torch
.load
(best_model_path
)
print
(f"成功加载模型: {best_model_path}"
)
char_to_idx = checkpoint['char_to_idx']
idx_to_char = checkpoint['idx_to_char']
vocab_size = len
(char_to_idx
)
model = RNN_Model
(
vocab_size=vocab_size,
embedding_dim=embedding_dim,
hidden_dim=hidden_dim,
num_layers=2
).to
(device
)
model
.load_state_dict
(checkpoint['model_state_dict']
)
start_
words = ["日", "红", "山", "夜", "湖", "君"]
generated_poems = []
for word in start_
words:
poem = generate_poem
(model, idx_to_char, char_to_idx, word
)
print
(f"\n--- Poem starting with '{word}' ---"
)
pretty_print_poem
(poem
)
generated_poems
.append
((word, poem
))
# 保存所有诗歌到文件
save_poems_to_file
(generated_poems
)
except Exception as e:
print
(f"\n加载模型失败: {e}"
)
print
("请检查模型文件是否存在或尝试重新训练模型"
)
输出
Using device: cuda
找到已有的训练模型，跳过训练
Generating poems with best model
...
成功加载模型:
./models/best_tangshi_model
.pth
加载模型失败: 'tuple' object has no attribute 'get'
请检查模型文件是否存在或尝试重新训练模型
进程已结束，退出代码为
0
解决上述问题
本文深入探讨了Java中Map接口的getOrDefault方法，该方法为获取映射值提供了便利，当键不存在时能返回默认值，避免了空指针异常，增强了代码的健壮性和可读性。
扫一扫
1354
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
