NLP（十七）：BiLSTM-Attention文本分类

 

一、概述

首先，我们导入需要的包，包括模型，优化器，梯度求导等，将数据类型全部转化成tensor类型

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable
import torch.nn.functional as F
import matplotlib.pyplot as plt

dtype = torch.FloatTensor

接下来我们确定一些基本的参数，并且简单地构造一个数据，实现情感的二分类。数据集中三个句子，一半正，一半负。label中1是好的情感，0是不好的情感。

embedding_dim = 3
n_hidden = 5 
num_classes = 2  # 0 or 1
sentences = ["i love you", "he loves me", "she likes baseball", "i hate you", "sorry for that", "this is awful"]
labels = [1, 1, 1, 0, 0, 0]  

接着，我们需要构建词表，把数据集中出现过的词拿出来并给它一个编号：

word_list = " ".join(sentences).split()
word_list = list(set(word_list))
word_dict = {w: i for i, w in enumerate(word_list)}
vocab_size = len(word_dict)

然后我们定义输入输出，输入其实是每个句子中的每个单词对应在词表中的id，将输入输出变成Variable，以便于求导：

inputs = []
for sen in sentences:
    inputs.append(np.asarray([word_dict[n] for n in sen.split()]))
targets = []
for out in labels:
    targets.append(out) 
input_batch = Variable(torch.LongTensor(inputs))
target_batch = Variable(torch.LongTensor(targets))

接下来构造模型：

class BiLSTM_Attention(nn.Module):
    def __init__(self):
        super(BiLSTM_Attention, self).__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        self.lstm = nn.LSTM(embedding_dim, n_hidden, bidirectional=True)
        self.out = nn.Linear(n_hidden * 2, num_classes)
    def attention_net(self, lstm_output, final_state):
        hidden = final_state.view(-1, n_hidden * 2, 1)  
        attn_weights = torch.bmm(lstm_output, hidden).squeeze(2) 
        soft_attn_weights = F.softmax(attn_weights, 1)
        context = torch.bmm(lstm_output.transpose(1, 2), soft_attn_weights.unsqueeze(2)).squeeze(2)
        return context, soft_attn_weights.data.numpy() 

    def forward(self, X):
        input = self.embedding(X) 
        input = input.permute(1, 0, 2)
        hidden_state = Variable(torch.zeros(1*2, len(X), n_hidden)) 
        cell_state = Variable(torch.zeros(1*2, len(X), n_hidden)) 
        output, (final_hidden_state, final_cell_state) = self.lstm(input, (hidden_state, cell_state))
        output = output.permute(1, 0, 2)
        attn_output, attention = self.attention_net(output, final_hidden_state)
        return self.out(attn_output), attention

首先embedding中需要传入词表，以及嵌入的维度。有一个双向LSTM层，还有一个线性层以获取LSTM中的隐层参数。
这里详细说一下attention层的操作，首先hidden 的维度是 [batch_size, n_hidden * num_directions(=2), 1(=n_layer)]，接下来确定attention矩阵，将LSTM输出与hidden相乘，去掉第三个维度。attn_weights的维度是[batch_size, n_step] ，两个矩阵相乘后的维度，[batch_size, n_hidden * num_directions(=2), n_step] * [batch_size, n_step, 1] = [batch_size, n_hidden * num_directions(=2), 1]，然后去掉了第三个维度的1。这样再经过softmax函数。再将权重函数与LSTM输出相乘得到context。最终context的维度就是 [batch_size, n_hidden * num_directions(=2)] 。
最后在forward方法中操作各个层，进行层的各种操作，获得输出和attention矩阵。

接下来就是将模型实例化，并确定损失函数，优化器：

model = BiLSTM_Attention()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

最后训练并测试：

# Training
for epoch in range(5000):
    optimizer.zero_grad()
    output, attention = model(input_batch)
    loss = criterion(output, target_batch)
    if (epoch + 1) % 1000 == 0:
        print('Epoch:', '%04d' % (epoch + 1), 'cost =', '{:.6f}'.format(loss))
    loss.backward()
    optimizer.step()
# Test
test_text = 'sorry hate you'
tests = [np.asarray([word_dict[n] for n in test_text.split()])]
test_batch = Variable(torch.LongTensor(tests))
# Predict
predict, _ = model(test_batch)
predict = predict.data.max(1, keepdim=True)[1]
if predict[0][0] == 0:
    print(test_text,"is Bad Mean...")
else:
    print(test_text,"is Good Mean!!")

最终我们画图看下attention中结果：

fig = plt.figure(figsize=(6, 3)) # [batch_size, n_step]
ax = fig.add_subplot(1, 1, 1)
ax.matshow(attention, cmap='viridis')
ax.set_xticklabels(['']+['first_word', 'second_word', 'third_word'], fontdict={'fontsize': 14}, rotation=90)
ax.set_yticklabels(['']+['batch_1', 'batch_2', 'batch_3', 'batch_4', 'batch_5', 'batch_6'], fontdict={'fontsize': 14})
plt.show()