诗歌生成比分类问题要稍微麻烦一些，而且第一次使用RNN做文本方面的问题，还是有很多概念性的东西~~~
数据下载：
链接：https://pan.baidu.com/s/1uCDup7U5rGuIlIb-lnZgjQ
提取码：f436
data.py——数据处理
import numpy as np
import os def get_data(conf):
'''
生成数据
:param conf: 配置选项,Config对象
:return: word2ix: 每个字符对应的索引id,如u'月'->100
:return: ix2word: 每个字符对应的索引id,如100->u'月'
:return: data: 每一行是一首诗对应的字的索引id
'''
if os.path.exists(conf.pickle_path): datas = np.load(conf.pickle_path) #np数据文件
data = datas['data']
ix2word = datas['ix2word'].item()
word2ix = datas['word2ix'].item()
return data, word2ix, ix2word
config.py——配置文件
class Config(object):
"""Base configuration class.For custom configurations, create a
sub-class that inherits from this one and override properties that
need to changed
""" # 模型保存路径前缀(几个epoch后保存)
model_prefix = 'checkpoints/tang' # 模型保存路径
model_path = 'checkpoints/tang.pth' # start words
start_words = '春江花月夜' # 生成诗歌的类型,默认为藏头诗
p_type = 'acrostic' # 训练次数
max_epech = 200 # 数据存放的路径
pickle_path = 'data/tang.npz' # mini批大小
batch_size =128###128 # dataloader加载数据使用多少进程
num_workers = 4 # LSTM的层数
num_layers = 2 # 词向量维数
embedding_dim = 128 # LSTM隐藏层维度
hidden_dim = 256 # 多少个epoch保存一次模型权重和诗
save_every = 10 # 训练是生成诗的保存路径
out_path = 'out' # 测试生成诗的保存路径
out_poetry_path = 'out/poetry.txt' # 生成诗的最大长度
max_gen_len = 200
use_gpu=True
model.py——模型
import torch.nn as nn
import torch
class PoetryModel(nn.Module):
def __init__(self, vocab_size, conf, device):
super(PoetryModel, self).__init__()
self.num_layers = conf.num_layers
self.hidden_dim = conf.hidden_dim
self.device = device
# 定义词向量层
self.embeddings = nn.Embedding(vocab_size, conf.embedding_dim)#（词库个数，词向量维度)
# 定义2层的LSTM，并且batch位于函数参数的第一位
self.lstm = nn.LSTM(conf.embedding_dim, conf.hidden_dim, num_layers=self.num_layers)
# 定义全连接层,后接一个softmax进行分类
self.linear_out = nn.Linear(conf.hidden_dim, vocab_size) def forward(self, input, hidden=None):
'''
:param input: (seq,batch)
:return: 模型的结果
'''
seq_len, batch_size = input.size()
embeds = self.embeddings(input) # embeds_size:(seq_len,batch_size,embedding_dim)
if hidden is None:
h0 = torch.zeros(self.num_layers, batch_size, self.hidden_dim).to(self.device)
c0 = torch.zeros(self.num_layers, batch_size, self.hidden_dim).to(self.device)
else:
h0, c0 = hidden
output, hidden = self.lstm(embeds, (h0, c0))#(seq_len,batch_size,隐藏层维度） output = self.linear_out(output.view(seq_len * batch_size, -1)) # output_size:(seq_len*batch_size,vocab_size)
return output, hidden
train.py——训练模型
import torch
from torch import nn
from torch.autograd import Variable
from torch.optim import Adam
from torch.utils.data import DataLoader from data import get_data
from model import PoetryModel
from config import Config
device=torch.device('cuda:0')
conf = Config() def generate(model, start_words, ix2word, word2ix, prefix_words=None):
'''
给定几个词，根据这几个词接着生成一首完整的诗歌
'''
print(start_words)
results = list(start_words)
start_word_len = len(start_words)
# 手动设置第一个词为<START>
# 这个地方有问题，最后需要再看一下
input = Variable(torch.Tensor([word2ix['<START>']]).view(1, 1).long())
if conf.use_gpu: input = input.cuda()
hidden = None if prefix_words:
for word in prefix_words:
output, hidden = model(input, hidden)
# 下边这句话是为了把input变成1*1？
input = Variable(input.data.new([word2ix[word]])).view(1, 1)
for i in range(conf.max_gen_len):
output, hidden = model(input, hidden)#input只有一个词，对应的是'<START>'的序号 if i < start_word_len:
w = results[i]
input = Variable(input.data.new([word2ix[w]])).view(1, 1)
else:
top_index = output.cpu().data.topk(1)[1][0].numpy()[0] w = ix2word[top_index]
results.append(w)
input = Variable(input.data.new([top_index])).view(1, 1)
if w == '<EOP>':
del results[-1] # -1的意思是倒数第一个
break
return results def gen_acrostic(model, start_words, ix2word, word2ix, prefix_words=None):
'''
生成藏头诗
start_words : u'深度学习'
生成：
深木通中岳，青苔半日脂。
度山分地险，逆浪到南巴。
学道兵犹毒，当时燕不移。
习根通古岸，开镜出清羸。
'''
results = []
start_word_len = len(start_words)
input = Variable(torch.Tensor([word2ix['<START>']]).view(1, 1).long())
if conf.use_gpu: input = input.cuda()
hidden = None index = 0 # 用来指示已经生成了多少句藏头诗
# 上一个词
pre_word = '<START>' if prefix_words:
for word in prefix_words:
output, hidden = model(input, hidden)
input = Variable(input.data.new([word2ix[word]])).view(1, 1) for i in range(conf.max_gen_len):
output, hidden = model(input, hidden)
top_index = output.data[0].topk(1)[1][0]
w = ix2word[top_index] if (pre_word in {u'。', u'！', '<START>'}):
# 如果遇到句号，藏头的词送进去生成 if index == start_word_len:
# 如果生成的诗歌已经包含全部藏头的词，则结束
break
else:
# 把藏头的词作为输入送入模型
w = start_words[index]
index += 1
input = Variable(input.data.new([word2ix[w]])).view(1, 1)
else:
# 否则的话，把上一次预测是词作为下一个词输入
input = Variable(input.data.new([word2ix[w]])).view(1, 1)
results.append(w)
pre_word = w
return results def train(**kwargs): for k, v in kwargs.items():
setattr(conf, k, v)
# 获取数据
data, word2ix, ix2word = get_data(conf)
# 生成dataload
dataloader = DataLoader(dataset=data, batch_size=conf.batch_size,
shuffle=True,
drop_last=True,
num_workers=conf.num_workers)
# 定义模型
model = PoetryModel(len(word2ix), conf, device).to(device)
# model.load_state_dict(torch.load(r'C:\Users\ocean\PycharmProjects\guesswhat_pytorch\checkpoints\tang_0.pth'))
# fout = open('%s/p%d' % (conf.out_path, 1), 'w',encoding='utf-8')
# # for word in list('春江花月夜'):
# # gen_poetry = generate(model, word, ix2word, word2ix)
# # fout.write("".join(gen_poetry) + "\n\n")
# gen_poetry = generate(model, list("北邮真的号"), ix2word, word2ix)
#
# fout.write("".join(gen_poetry) + "\n\n")
# fout.close()
# torch.save(model.state_dict(), "%s_%d.pth" % (conf.model_prefix, 1)) # 定义优化器
optimizer = Adam(model.parameters())
# 定义损失函数
criterion = nn.CrossEntropyLoss()
# 开始训练模型
for epoch in range(conf.max_epech):
epoch_loss = 0
for i, data in enumerate(dataloader): data = data.long().transpose(1, 0).contiguous()#(sequence长度，batch_size) input, target = data[:-1, :], data[1:, :]
input, target = input.to(device), target.to(device)
optimizer.zero_grad()
output, _ = model(input) loss = criterion(output, target.view(-1)) loss.backward()
optimizer.step()
epoch_loss += loss.item()
print("epoch_%d_loss:%0.4f" % (epoch, epoch_loss / conf.batch_size))
if epoch % conf.save_every == 0:
fout = open('%s/p%d' % (conf.out_path, epoch), 'w',encoding='utf-8')
for word in list('春江花月夜'):
gen_poetry = generate(model, word, ix2word, word2ix)
fout.write("".join(gen_poetry) + "\n\n")
fout.close()
torch.save(model.state_dict(), "%s_%d.pth" % (conf.model_prefix, epoch)) if __name__ == '__main__': train()
最终效果：
春雨，君王背日暮花枝。桂花飘雨裛芙蓉，花蕚垂红绾芙蓉。上天高峨落不归，中有一枝春未老。一枝香蘂红妆结，春风吹花飘落萼。今朝今日凌风沙，今日还家花落早。东风吹落柳条生，柳色参差春水东。昨日风烟花满树，今日东风正如萍。杏园春色不自胜，青春忽倒春风来。春风不哢花枝落，况复春风花满枝。 江上春未央，春光照四面。一日一千里，一朝一瞬息。不如塌然云，不见巢下树。一身一何讬，万事皆有敌。君子不敢横，君心若为役。呜呼两鬓苦，又如寒玉翦。不知何代费，所以心不殒。一朝得之愚，所以心所施。我亦我未领，我来亦未归。始知与君子，不觉身不饥。彭泽有余事，吾君何所为。何以为我人，於今有耆谁。 花间一人家，十五日中见。一朝出门门，不见君子诺。车骑徒自媒，朱绂不能竞。拜军拜车骑，倏忽嫖姚羌。既无征镇愤，慷慨望乡国。一朝辞虏府，暮宿在蓟垒。君子傥封侯，今人在咸朔。英英在其间，日昃不敢作。云山互相见，魏阙空怀戚。何必在沛人，裴回眇眇。所念无穷，斯人不怠。。 月白风来吹，君心不可攀。从来一字内，不觉一朝闲。未达身难弃，衰容事不闲。不忧讥孺子，不觉老农闲。寝食能供给，闲橙媿漉肱。酒阑湘口臥，窗拔峡添灯。静谭畬茶骇，遥闻夜笛闲。芦洲多雨霁，石火带霜蒸。酿酒眠新熟，扁舟醉自闲。夜渔疎竹坞，春水钓渔关。石笋穿云烧，江花带笋斑。此时多好客，不敢问山僧。 夜夜拍人笑，春风弄酒丝。花开桃李岭，花落洞庭春。酒思同君醉，诗成是袜尘。自怜心已矣，何事梦何如。摈世才难易，伤心镜不如。脸如银烛薄，色映玉楼嚬。绣户雕筵软，鸳鸯拂枕春。相逢期洛下，梦想忆扬秦。玉匣调金鼎，金盘染髻巾。鷰人曾有什，山寺不相亲。鹤毳应相毒，蝇蚊爽有真。空余襟袖下，不觉世间人。
参考博客：https://blog.csdn.net/jiangpeng59/article/details/81003058