Pytorch基于 LSTM 的自动写诗

实验环境

Pytorch 1.4.0
conda 4.7.12
Jupyter Notebook 6.0.1
Python 3.7

数据集介绍

chinese-poetry: 最全中文诗歌古典文集数据库
最全的中华古典文集数据库，包含 5.5 万首唐诗、26 万首宋诗、2.1 万首宋词和其他古典文集。诗人包括唐宋两朝近 1.4 万古诗人，和两宋时期 1.5 千古词人。数据来源于互联网。
实验使用预处理过的二进制文件 tang.npz 作为数据集，含有 57580 首唐诗，每首诗限定在 125 词，不足 125 词的以填充。数据集以 npz 文件形式保存，包含三个部分:

data: (57580,125) 的 numpy 数组，总共有 57580 首诗歌，每首诗歌长度为 125 字符 (不足 125 补空格，超过 125 的丢弃)，将诗词中的字转化为其在字典中的序号表示
ix2word: 序号到字的映射，每个序号和它对应的词，例如序号 898 对应“雪”
word2ix: 字到序号的映射，每个字和它对应的序号，例如“雪”对应序号是 898

训练过程

数据准备

首先，导入实验所需的库，定义一些宏参数，BATCH_SIZE 表示每个 batch 加载多少个样本、 EPOCHS 表示总共训练批次。如果支持 cuda 就用 gpu 来 run，不支持就用 cpu 来 run。

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import time
import matplotlib.pyplot as plt
from torch.optim.lr_scheduler import *
from torch.utils.data import DataLoader
from torch.autograd import Variable

BATCH_SIZE = 16
EPOCHS = 4 
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

我先读取预处理好的数据，再将其转为 Tensor，这次没有使用 Dataset 去封装 data，但是它还是可以利用 Dataloader 进行多线程加载，因为 data 作为一个 Tensor 对象，自身已经实现了 getitem 和 len 方法。接着构造数据提供器，shuffle=True 设置在每个 epoch 重新打乱数据，保证数据的随机性;还有多线程加载数据。处理好了训练数据。

# 读入预处理的数据
datas = np.load("./tang.npz")
data = datas['data']
ix2word = datas['ix2word'].item()
word2ix = datas['word2ix'].item()
    
# 转为torch.Tensor
data = torch.from_numpy(data)
train_loader = DataLoader(data, batch_size = BATCH_SIZE, shuffle = True, num_workers = 2)

网络配置

定义一个循环神经网络，输入的字词序号经过 nn.Embedding 得到相应词的词向量表示，然后利用 3 层 LSTM 提取词的所有隐藏元信息，再利用隐藏元的信息进行分类，判断输出属于每一个词的概率。然后通过全连接的输出层将词向量升维回字词序号，全连接输出层所有激活函数都使用 ReLU 函数。

class PoetryModel(nn.Module):
    def __init__(self, vocab_size, embedding_dim, hidden_dim):
        super(PoetryModel, self).__init__()
        self.hidden_dim = hidden_dim
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        self.lstm = nn.LSTM(embedding_dim, self.hidden_dim, num_layers=3)
        self.classifier=nn.Sequential(
            nn.Linear(self.hidden_dim, 512), 
            nn.ReLU(inplace=True), 
            nn.Linear(512, 2048), 
            nn.ReLU(inplace=True),
            nn.Linear(2048, vocab_size)
        )

    def forward(self, input, hidden = None):
        seq_len, batch_size = input.size()
        
        if hidden is None:
            h_0 = input.data.new(3, batch_size, self.hidden_dim).fill_(0).float()
            c_0 = input.data.new(3, batch_size, self.hidden_dim).fill_(0).float()
        else:
            h_0, c_0 = hidden

        embeds = self.embedding(input)
        output, hidden = self.lstm(embeds, (h_0, c_0))
        output = self.classifier(output.view(seq_len * batch_size, -1))
        
        return output, hidden

初始化模型;如果继续上一次的训练，就设置预训练模型路径配置模型;优化器采用 Adam;损失
函数使用交叉熵损失;使用 StepLR 在每 4 个 epoch 结束时调整学习率。

# 配置模型，是否继续上一次的训练
model = PoetryModel(len(word2ix),embedding_dim = 128,hidden_dim = 256)

model_path = ''         # 预训练模型路径
if model_path:
    model.load_state_dict(torch.load(model_path))
model.to(DEVICE)
    
optimizer = torch.optim.Adam(model.parameters(), lr = 5e-3)
# optimizer = torch.optim.SGD(model.parameters(), lr=5e-3, momentum=0.9, weight_decay=5e-4)
criterion = nn.CrossEntropyLoss()
scheduler = StepLR(optimizer, step_size=10)

模型训练


def train(model, dataloader, ix2word, word2ix, device, optimizer, scheduler, epoch):
    model.train()
    train_loss = 0.0
    
    for batch_idx, data in enumerate(dataloader):
        data = data.long().transpose(1, 0).contiguous()
        data = data.to(device)
        optimizer.zero_grad()
        input, target = data[:-1, :], data[1:, :]
        output, _ = model(input)
        loss = criterion(output, target.view(-1))
        loss.backward()  
        optimizer.step()
        train_loss += loss.item()
            
        if (batch_idx+1) % 200 == 0:
            print('train epoch: {} [{}/{} ({:.0f}%)]\tloss: {:.6f}'.format(
                epoch, batch_idx * len(data[1]), len(dataloader.dataset),
                100. * batch_idx / len(dataloader), loss.item()))
            
    train_loss *= BATCH_SIZE
    train_loss /= len(train_loader.dataset)
    print('\ntrain epoch: {}\t average loss: {:.6f}\n'.format(epoch,train_loss))
    scheduler.step()
    
    return train_loss

train_losses = []

for epoch in range(1,EPOCHS+1):
    tr_loss = train(model,train_loader,ix2word,word2ix,DEVICE,optimizer,scheduler,epoch)
    train_losses.append(tr_loss)
    
# 保存模型
filename = "model" + str(time.time()) + ".pth"
torch.save(model.state_dict(), filename)

模型评估

定义给定诗歌的开头几个字接着生成整首诗歌的函数，生成步骤如下:首先利用给定的文字计算隐藏元，并预测下一个词;将上一步计算的隐藏元和输出作为下一步的输入继续预测新的输出和计算隐藏元;prefix_word 可以用来控制生成诗歌的意境和长短。

def generate(model, start_words, ix2word, word2ix, max_gen_len, prefix_words=None):
    # 读取唐诗的第一句
    results = list(start_words)
    start_word_len = len(start_words)
    
    # 设置第一个词为<START>
    input = torch.Tensor([word2ix['<START>']]).view(1, 1).long()
    input = input.to(DEVICE)
    hidden = None
    
    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(max_gen_len):
        output, hidden = model(input, hidden)
        # 读取第一句
        if i < start_word_len:
            w = results[i]
            input = input.data.new([word2ix[w]]).view(1, 1)
        # 生成后面的句子
        else:
            top_index = output.data[0].topk(1)[1][0].item()
            w = ix2word[top_index]
            results.append(w)
            input = input.data.new([top_index]).view(1, 1)
        # 结束标志
        if w == '<EOP>':
            del results[-1]
            break
            
    return results

start_words 是唐诗的第一句，max_gen_len 是生成唐诗的最长长度。

start_words = '白日依山尽'  # 唐诗的第一句
max_gen_len = 128        # 生成唐诗的最长长度

prefix_words = None
results = generate(model, start_words, ix2word, word2ix, max_gen_len, prefix_words)
poetry = ''
for word in results:
    poetry += word
    if word == '。' or word == '!':
        poetry += '\n'
        
print(poetry)

生成藏头诗的步骤如下:

先输入要藏头的字，开始预测下一个字;
上一步预测的字作为输入，继续预测下一个字。
重复第二步，直到输出的字是“。”或者“!”，说明一句诗结束了，可以继续输入下一句藏头的
字，跳到第一步。
重复上述步骤直到所有藏头的字都输入完毕。

def gen_acrostic(model, start_words, ix2word, word2ix, prefix_words=None):
    # 读取唐诗的“头”
    results = []
    start_word_len = len(start_words)
    
    # 设置第一个词为<START>
    input = (torch.Tensor([word2ix['<START>']]).view(1, 1).long())
    input = input.to(DEVICE)
    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(max_gen_len_acrostic):
        output, hidden = model(input, hidden)
        top_index = output.data[0].topk(1)[1][0].item()
        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 = (input.data.new([word2ix[w]])).view(1, 1)
                
        # 否则，把上一次预测作为下一个词输入
        else:
            input = (input.data.new([word2ix[w]])).view(1, 1)
        results.append(w)
        pre_word = w
        
    return results

start_words_acrostic = '深度学习'  # 唐诗的“头”
max_gen_len_acrostic = 120               # 生成唐诗的最长长度
prefix_words = None
results_acrostic = gen_acrostic(model, start_words_acrostic, ix2word, word2ix,prefix_words)

poetry = ''
for word in results_acrostic:
    poetry += word
    if word == '。' or word == '!':
        poetry += '\n'
        
print(poetry)

结果

生成的诗歌有些质量较好，学会了一些简单的对偶和押韵。但是随着生成的诗歌的长度变长，诗歌的意境会变化，后面的和前面的相差甚远