循环神经网络小项目 | 七绝作诗

数据集

做深度的第一步就是要收集炼丹原料准备数据集，这里我选择了 GitHub 上的这个项目 chinese-poetry 然后取其中的唐诗部分，并去除含有现在字库缺少的字的唐诗，将诗中一些注解删除（比如一字多解），再利用正则匹配提取出七绝唐诗，将每个汉字拆成单独的词，在唐诗开头和结尾添加 <SOP> 和 <EOP> 符号标志开始和结束，一共准备了 10922 首七绝唐诗作为训练数据。

数据集长这样：

<SOP> 中 管 五 弦 初 半 曲 ， 遙 教 合 上 隔 簾 聽 。 一 聲 聲 向 天 頭 落 ， 效 得 仙 人 夜 唱 經 。 <EOP>
<SOP> 自 直 梨 園 得 出 稀 ， 更 番 上 曲 不 教 歸 。 一 時 跪 拜 霓 裳 徹 ， 立 地 階 前 賜 紫 衣 。 <EOP>
<SOP> 旋 翻 新 譜 聲 初 足 ， 除 却 梨 園 未 教 人 。 宣 與 書 家 分 手 寫 ， 中 官 走 馬 賜 功 臣 。 <EOP>
<SOP> 伴 教 霓 裳 有 貴 妃 ， 從 初 直 到 曲 成 時 。 日 長 耳 裏 聞 聲 熟 ， 拍 數 分 毫 錯 總 知 。 <EOP>
<SOP> 弦 索 摐 摐 隔 綵 雲 ， 五 更 初 發 一 山 聞 。 武 皇 自 送 西 王 母 ， 新 換 霓 裳 月 色 裙 。 <EOP>

模型

模型搭建非常简单，就是普通的双层单向 LSTM。每一个汉字首先经过 Embedding 层，然后输入 LSTM，最后将 LSTM 的输出经过一个线性层解码得到输出的词。

class BasicRNN(nn.Module):
    def __init__(self, vocab_size, embedding_dim, hidden_dim, num_layers, dropout=0.5):
        super(BasicRNN, self).__init__()
        self.vocab_size = vocab_size
        self.embedding_dim = embedding_dim
        self.embed = nn.Embedding(vocab_size, embedding_dim)
        self.lstm = nn.LSTM(embedding_dim, hidden_dim, num_layers, dropout=dropout)
        self.decode = nn.Linear(hidden_dim, vocab_size)
        self.dropout = nn.Dropout(dropout)
        self.num_layers = num_layers
        self.hidden_dim = hidden_dim
        self.init_weights()
        
    def forward(self, x, hidden_states):
        embedding = self.dropout(self.embed(x))
        output, (h_n, c_n) = self.lstm(embedding, hidden_states)
        output = self.dropout(output)
        bsz = output.size(1)
        decoded = self.decode(output.view(-1, output.size(2)))
        decoded = decoded.view(-1, bsz, self.vocab_size)
        return decoded, (h_n, c_n)
    
    def init_weights(self):
        initrange = 0.1
        self.embed.weight.data.uniform_(-initrange, initrange)
        self.decode.bias.data.zero_()
        self.decode.weight.data.uniform_(-initrange, initrange)
    
    def init_hidden(self, bsz):
        weight = next(self.parameters())
        return (weight.new_zeros(self.num_layers, bsz, self.hidden_dim),
                weight.new_zeros(self.num_layers, bsz, self.hidden_dim))

需要注意的是输出层不需要任何的激活函数，这里当时调试了好久，loss 无论如何都下降不了，后来发现是 nn.CrossEntropyLoss 会自带一个 softmax 激活函数。

训练

实际上这个作诗模型是一个语言模型（Language Model），为了简化操作，我用了 torchtext 中的 BPTTIterator 来生成 Mini Batch。

需要注意的是，隐藏层每次都需要和之前的历史记录分离开来，否则梯度会一直回传下去。

def train(model, dataset, lr=1e-3, epochs=10, start=0, save_per=1000, debug=False):
    train_iter = torchtext.data.BPTTIterator(
            dataset,
            batch_size=2048,
            bptt_len=33,
            device=device,
            repeat=False
        )
    vocab_size = len(dataset.fields['text'].vocab)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=lr)
    hidden = None
    total_loss = []
    for epoch in range(epochs):
        epoch = epoch + start
        try:
            model.train()
            epoch_loss = []
            train_iter.init_epoch()
            for i, batch in enumerate(tqdm(train_iter)):
                if hidden is None:
                    hidden = model.init_hidden(batch.batch_size)
                else:
                    hidden = detach_hidden(hidden)

                text, target = batch.text, batch.target
                output, hidden = model(text, hidden)
                optimizer.zero_grad()
                loss = criterion(output.view(-1, vocab_size), target.view(-1))
                loss.backward()
                optimizer.step()
                epoch_loss.append(loss.item())

            epoch_loss = np.mean(epoch_loss)
            total_loss.append(epoch_loss)
            if debug:
                print("Epoch %d Loss: %f" % (epoch, epoch_loss))
                print(''.join(generate_poem(model)))
            elif (epoch + 1) % 10 == 0: 
                print("Epoch %d Loss: %f" % (epoch, epoch_loss))
                print(''.join(generate_poem(model)))
                print(''.join(generate_poem(model, True)))
                with open("loss.log", "a") as f:
                    f.write("Epoch %d Loss: %f\n" % (epoch, epoch_loss))
                    f.write(''.join(generate_poem(model)) + '\n')
                    f.write(''.join(generate_poem(model, True)) + '\n')
            if (epoch + 1) % save_per == 0 or epoch == 0 and not debug:
                torch.save(model.state_dict(), "model_{0:d}.pth".format(epoch))
        except KeyboardInterrupt:
            torch.save(model.state_dict(), "model_{0:d}.pth".format(epoch))
            return total_loss
    return total_loss

我采用的训练配置是：

batch_size=2048
BPTT=33
embedding_size=300
num_layers=2
hidden_dim=1024

该配置在 NVIDIA GTX Titan X 上消耗大约 11GB 显存，约 10 秒运行一个 epoch，最终花了两天时间训练了 15755 个 epoch。

生成

生成唐诗则是传进去一个 <SOP> 标识符，然后将上一时刻的输出作为下一时刻的输入，直到遇到 <EOP> 则停止生成。

从输出中生成古诗有两种办法，一种是直接取概率最大的字作为输出，另一种则是按照概率随机采样。

def generate_poem(model, sample=False):
    model.eval()
    idx = TEXT.vocab.stoi["<SOP>"]
    x = torch.Tensor([idx]).view(1, 1).long().to(device)
    poem = []
    hidden = model.init_hidden(1)
    with torch.no_grad():
        for _ in range(128):
            output, hidden = model(x, hidden)
            output = output.view(model.vocab_size)
            if sample:
                probs = F.softmax(output, dim=0).cpu().numpy()
                probs /= probs.sum()
                idx = np.random.choice(range(model.vocab_size), p=probs)
            else:
                idx = torch.argmax(output)
            if idx == TEXT.vocab.stoi["<EOP>"]: break
            poem.append(TEXT.vocab.itos[idx])
            x = torch.Tensor([idx]).view(1, 1).long().to(device)
    return poem

最终生成的古诗效果如下：

Epoch 15729 Loss: 1.304158
君不到山無處物，一生無事與身閑。人間盡說逢花雨，不是一生一覺塵。
如中百歲曾留得，遙看還鄉夢覺看。不知日夜東山遠，獨照紅塵滿地花。
Epoch 15739 Loss: 1.304422
興不見君心不知，空留一鶴到山邊。莫言花重船應沒，自有人間不不知。
看花莫羨新條在，花裏人呼萬古同。聞道不堪猶自異，兩頭分上一枝梅。
Epoch 15749 Loss: 1.305263
興不見君來未歸，今朝同向五湖中。相逢一宿最高寺，半夜不知何處去。
今年閑向人中見，已見臨川五月月。青山山下何日期，西林宿竹獨相思。

完整代码和预训练的模型：https://github.com/howardlau1999/char-rnn-poet

改进

加 Attention
加平仄信息
改成 VAE 模型