深度学习实战（基于pytroch）系列（四十一）长短期记忆（LSTM）pytorch简洁实现

长短期记忆（LSTM）pytorch简洁实现

上一节我们已经"从零开始实现了LSTM"，本节将使用PyTorch来更简洁地实现长短期记忆（LSTM）语言模型。

首先，我们读取周杰伦专辑歌词数据集。这一步和上节的代码基本一致。

python 复制代码

import torch
import torch.nn as nn
import torch.optim as optim
import math
import random
import time
import zipfile

def load_data_jay_lyrics():
    with zipfile.ZipFile('../data/jaychou_lyrics.txt.zip') as zin:
        with zin.open('jaychou_lyrics.txt') as f:
            corpus_chars = f.read().decode('utf-8')
    corpus_chars = corpus_chars.replace('\n', ' ').replace('\r', ' ')
    corpus_chars = corpus_chars[0:10000]
    idx_to_char = list(set(corpus_chars))
    char_to_idx = {char: i for i, char in enumerate(idx_to_char)}
    vocab_size = len(char_to_idx)
    corpus_indices = [char_to_idx[char] for char in corpus_chars]
    return corpus_indices, char_to_idx, idx_to_char, vocab_size


(corpus_indices, char_to_idx, idx_to_char, vocab_size) = load_data_jay_lyrics()

定义模型

PyTorch的nn模块提供了循环神经网络的实现。下面构造一个含单隐藏层、隐藏单元个数为256的循环神经网络层_layer 。

python 复制代码

num_hiddens = 256

lstm_layer = nn.LSTM(input_size=vocab_size, hidden_size=num_hiddens)

接下来我们初始化隐藏状态。这里和我前面学的rnn和gru模型不同，因为LSTM需要两个状态：隐藏状态(hidden state)和细胞状态(cell state)。

python 复制代码

batch_size = 2
num_steps = 35

h_0 = torch.zeros((1, batch_size, num_hiddens))  # 隐藏状态
c_0 = torch.zeros((1, batch_size, num_hiddens))  # 细胞状态
state = (h_0, c_0)

下面的代码我们可以看到经过lstm_layer层后得到的输出形状。

python 复制代码

X = torch.rand(num_steps, batch_size, vocab_size)
Y, state_new = lstm_layer(X, state)
Y.shape

输出

torch.Size([35, 2, 256])

接下来我们继承Module类来定义一个完整的循环神经网络。它首先将输入数据使用one-hot向量表示后输入到gru_layer 中，然后使用全连接输出层得到输出。输出个数等于词典大小vocab_size。这里在begin_state会进入到lstm分支中。

python 复制代码

from torch.nn import functional as F
class RNNModel(nn.Module):
    def __init__(self, rnn_layer, vocab_size, **kwargs):
        super(RNNModel, self).__init__(**kwargs)
        self.rnn = rnn_layer
        self.vocab_size = vocab_size
        self.dense = nn.Linear(num_hiddens, vocab_size)

    def forward(self, inputs, state):
        # 将输入转置成(num_steps, batch_size)后获取one-hot向量表示
        X = F.one_hot(inputs.T.long(), self.vocab_size)
        X = X.to(torch.float32)
        Y, state = self.rnn(X, state)
        # 全连接层会首先将Y的形状变成(num_steps * batch_size, num_hiddens)，它的输出
        # 形状为(num_steps * batch_size, vocab_size)
        output = self.dense(Y.reshape(-1, Y.shape[-1]))
        return output, state

    def begin_state(self, batch_size, device):
        if not isinstance(self.rnn, nn.LSTM):
            # nn.RNN以张量作为隐藏状态
            return torch.zeros((self.rnn.num_layers, batch_size, self.rnn.hidden_size), device=device)
        else:
            # nn.LSTM以元组作为隐藏状态
            return (torch.zeros((self.rnn.num_layers, batch_size, self.rnn.hidden_size), device=device),
                    torch.zeros((self.rnn.num_layers, batch_size, self.rnn.hidden_size), device=device))

训练模型

同之前一样，下面定义一个预测函数。这里的实现区别在于前向计算和初始化隐藏状态的函数接口。

python 复制代码

def predict_rnn_pytorch(prefix, num_chars, model, vocab_size, device, idx_to_char,
                      char_to_idx):
    # 使用model的成员函数来初始化隐藏状态
    state = model.begin_state(batch_size=1, device=device)
    output = [char_to_idx[prefix[0]]]
    for t in range(num_chars + len(prefix) - 1):
        X = torch.tensor([output[-1]], device=device).reshape((1, 1))
        (Y, state) = model(X, state)  # 前向计算不需要传入模型参数
        if t < len(prefix) - 1:
            output.append(char_to_idx[prefix[t + 1]])
        else:
            output.append(int(Y.argmax(dim=1).item()))
    return ''.join([idx_to_char[i] for i in output])

让我们使用权重为随机值的模型来预测一次。

python 复制代码

def try_gpu():
    """If GPU is available, return torch.device('cuda'); else return torch.device('cpu')."""
    if torch.cuda.is_available():
        return torch.device('cuda')
    else:
        return torch.device('cpu')

device = try_gpu()
model = RNNModel(gru_layer, vocab_size)
model = model.to(device)
predict_rnn_pytorch('分开', 10, model, vocab_size, device, idx_to_char, char_to_idx)

输出

'分开瞎瞎行行极极行枚极极'

接下来实现训练函数。算法同之前的一样。

python 复制代码

def grad_clipping(params, theta, device):
    norm = torch.tensor([0.0], device=device)
    for param in params:
        norm += (param.grad.data ** 2).sum()
    norm = norm.sqrt().item()
    if norm > theta:
        for param in params:
            param.grad.data *= (theta / norm)


def data_iter_consecutive(corpus_indices, batch_size, num_steps, device=None):
    corpus_indices = torch.tensor(corpus_indices, device=device)
    data_len = len(corpus_indices)
    batch_len = data_len // batch_size
    indices = corpus_indices[0: batch_size * batch_len].reshape(
        batch_size, batch_len)
    epoch_size = (batch_len - 1) // num_steps
    for i in range(epoch_size):
        i = i * num_steps
        X = indices[:, i: i + num_steps]
        Y = indices[:, i + 1: i + num_steps + 1]
        yield X, Y


def train_and_predict_rnn_pytorch(model, num_hiddens, vocab_size, device,
                                corpus_indices, idx_to_char, char_to_idx,
                                num_epochs, num_steps, lr, clipping_theta,
                                batch_size, pred_period, pred_len, prefixes):
    loss = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=lr)
    
    for epoch in range(num_epochs):
        l_sum, n, start = 0.0, 0, time.time()
        data_iter = data_iter_consecutive(
            corpus_indices, batch_size, num_steps, device)
        state = model.begin_state(batch_size=batch_size, device=device)
        for X, Y in data_iter:
            if isinstance(state, tuple):  # LSTM, state:(h, c)  
                state = (state[0].detach(), state[1].detach())
            else:
                state = state.detach()
            
            (output, state) = model(X, state)
            y = Y.T.reshape(-1)
            l = loss(output, y.long())
            
            optimizer.zero_grad()
            l.backward()
            # 梯度裁剪
            grad_clipping(model.parameters(), clipping_theta, device)
            optimizer.step()
            l_sum += l.item() * y.numel()
            n += y.numel()

        if (epoch + 1) % pred_period == 0:
            print('epoch %d, perplexity %f, time %.2f sec' % (
                epoch + 1, math.exp(l_sum / n), time.time() - start))
            for prefix in prefixes:
                print(' -', predict_rnn_pytorch(
                    prefix, pred_len, model, vocab_size, device, idx_to_char,
                    char_to_idx))

使用和之前实验中一样的超参数来训练模型。

python 复制代码

num_epochs, batch_size, lr, clipping_theta = 250, 32, 1e-3, 1e-2
pred_period, pred_len, prefixes = 50, 50, ['分开', '不分开']
train_and_predict_rnn_pytorch(model, num_hiddens, vocab_size, device,
                            corpus_indices, idx_to_char, char_to_idx,
                            num_epochs, num_steps, lr, clipping_theta,
                            batch_size, pred_period, pred_len, prefixes)

输出

epoch 50, perplexity 20.578601, time 0.03 sec

分开我不不你不要再我我不不不不不我不不不我不不不不不再我不不再我不不再我不不不

不分开我不不不知道我想要你我不不我不不不不不再我不不再我不不不我不不不我不不不我不

epoch 100, perplexity 2.377903, time 0.03 sec

分开我的可爱你我的没有你你的没有有一直对不我想要你你的我不不要再我我不要你的我不

不分开我不起你知道我想要你汉堡我想要你这样打我妈妈我的你爱你我不要再想要我不要我不不

epoch 150, perplexity 1.290544, time 0.03 sec

分开我的可却为人看着我别发抖快给我抬起头有话去对医药箱说别怪我别怪我说你怎么面对我甩

不分开我想要你想要再想要我想要你想要我想回回忆爱上的让我说说你什么我不懂说你了其实我的

epoch 200, perplexity 1.089689, time 0.03 sec

分开我的可你为什么我打我爸爸你的那你怎么我想要和你的宙去想要和你融化在一起融化在宇宙里

不分开我的可不要再想我想要你我的回不要再这样打的手我的可爱你你想要我想想你想要你对我

epoch 250, perplexity 1.051081, time 0.03 sec

分开我的感动一场梦瞎透了一定一个人轻就是跟着我一口在人回著一起一的落我们在睡着

不分开我不开不知不觉我跟了这节奏后知后觉又过了一个秋后知后觉我该好好生活我该好好生活不知

本系列目录链接