深度学习实战（基于pytroch）系列（三十九）门控循环单元（GRU）pytorch简洁实现

门控循环单元（GRU）pytorch简洁实现

上一节我们已经"从零开始实现了GRU"，本节将使用PyTorch来更简洁地实现基于门控循环单元（GRU）的语言模型。

这节的代码本质和前面"循环神经网络的pytorch简洁实现"差不多，只是将rnn_layer变成了gru_layer。

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

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的循环神经网络层gru_layer 。

python 复制代码

num_hiddens = 256
gru_layer = nn.GRU(input_size=vocab_size, hidden_size=num_hiddens)

接下来我们初始化隐藏状态。它有一个形状为(层数, 批量大小, 隐藏单元个数)的元素。

python 复制代码

batch_size = 2
state = torch.zeros((1, batch_size, num_hiddens))
state.shape

输出

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

与之前实现的循环神经网络不同，这里gru_layer 的输入形状为(时间步数, 批量大小, 输入维度)。其中输入维度即one-hot向量长度（词典大小）。此外，gru_layer 在前向计算后会分别返回输出和隐藏状态，其中输出指的是隐藏层在各个时间步上计算并输出的隐藏状态，它们通常作为后续输出层的输入。需要强调的是，该"输出"本身并不涉及输出层计算，形状为(时间步数, 批量大小, 隐藏单元个数)。而RNN实例在前向计算返回的隐藏状态指的是隐藏层在最后时间步的可用于初始化下一时间步的隐藏状态。

python 复制代码

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

输出

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

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

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 5.026790, time 0.02 sec

分开始共书书书书书书书书书书书书书书书书书书书书书书书书书书书书书书书书书书书书书书书书书书书书书书书书

不分开始打我我不能再想我不能再想我不能再想我不能再想我不能再想我不能再想我不能再想我不能再

epoch 100, perplexity 1.194557, time 0.02 sec

分开始楷楷书杵杵杵杵杵杵杵杵杵杵杵杵杵杵杵杵杵杵杵杵杵杵杵杵杵杵杵杵杵杵杵杵杵杵杵印印印印印印印印印印印

不分开始打我妈妈难道你手不会痛吗不要再这样打我妈妈难道你手不会痛吗我叫你怎么面对我甩开球我满腔的

epoch 150, perplexity 1.049898, time 0.02 sec

分开始所所有话所所所所所所所所所所所所所所所所所所所所所所所所所所所所所所所所所所所所所所所所所所所所所

不分开始打我相思寄红豆相思寄红豆无能为力的在人海中漂泊心伤透娘子她人在江南等我泪不休语沉默娘子

epoch 200, perplexity 1.024190, time 0.02 sec

分开平常话话常话话话话话话话话话话话话话话话话话话话话话话话话话话话话话话话话话话话话话话话话话话话话话

不分开始打我我打我妈妈难道你手不会痛吗其实我回家就想要阻止一切让家庭回到过去甜甜温馨的欢乐香味

epoch 250, perplexity 1.016784, time 0.02 sec

分开平常话话去就因为喝醉酒他就能拿我妈出气我真的看不下去以为我较细汉从小到大只有妈妈的温暖

不分开始打我我跟了这样担忧唱着歌一直走我想就这样牵着你的手不放开爱可不可以简简单单没有伤害你

本系列目录链接