Seq2Seq+Attention代码

题目

'''
Description: Seq2Seq+Attention
Autor: 365JHWZGo
Date: 2021-12-18 18:24:19
LastEditors: 365JHWZGo
LastEditTime: 2021-12-26 16:57:05
'''

翻译任务

将Chinese time翻译成English time

Chinese time order: yy/mm/dd [‘31-04-26’, ‘04-07-18’, ‘33-06-06’]
English time order: dd/M/yyyy [‘26/Apr/2031’, ‘18/Jul/2004’, ‘06/Jun/2033’]

可选择词汇，共27个
enc_v_dim = 27
dec_v_dim = 27
Vocabularies: {‘7’, ‘0’, ‘-’, ‘9’, ‘Apr’, ‘4’, ‘Jun’, ‘Feb’, ‘’, ‘’, ‘Jan’, ‘May’, ‘Oct’, ‘/’, ‘3’, ‘Nov’, ‘Jul’, ‘Aug’, ‘Mar’, ‘2’, ‘6’, ‘Sep’, ‘Dec’, ‘’, ‘8’, ‘1’, ‘5’}

x示例输入和张量表示
x index sample:
31-04-26
[6 4 1 3 7 1 5 9]

y示例输入和张量表示
y index sample:
26/Apr/2031
[14 5 9 2 15 2 5 3 6 4 13]

比较和单纯的Seq2Seq的区别

加入注意力机制后就是使得每次的LSTM的h，c都是集中于翻译相关值生成的。

注意力的计算用到了General Attention，它的得分公式是
$$Score(k,q)=k^T{W}_{a}q$$

这里的 q 是一个进行注意的对象，在任务中为decoder中新生成的hx（隐状态）
$k$是被注意的对象，在任务中为encoder中的输出o

意思是将decoder每预测一个词， 我都拿着这个decoder现在的信息q去和encoder输出的所有信息k做注意力的计算。

注意力分配过程

# hx 相当于q，decoder生成的hidden state
# o 相当于k，encoder所有信息
# w*q
self.attn(hx.unsqueeze(1) 
# k^t*w*q
score = torch.matmul(self.attn(hx.unsqueeze(1)),o.permute(0,2,1))
# 将得分归一化
score01 = softmax(score, dim=2)
# 不同的权重代表不同的注意力分配，在decoder输出hx的作用下
attnDistribute = torch.matmul(score01,o)

根据输入得到预测的hx

# 计算得到的hx
hx, cx = self.decoder_cell(dec_in, (hx, cx))

预测值的hx和注意力分配相结合

ha = torch.cat([context.squeeze(1),hx],dim=1) 

得到预测的值

result = self.decoder_dense(ha)   

下图是一个简单举例

完整代码

from torch import nn
import torch
import utils
from torch.utils.data import DataLoader
from torch.nn.functional import cross_entropy,softmax

class Seq2Seq(nn.Module):
    def __init__(self,enc_v_dim, dec_v_dim, emb_dim, hidden_size, max_pred_len, start_token, end_token):
        super().__init__()
        self.hidden_size = hidden_size                  # hidden_size
        self.dec_v_dim = dec_v_dim
        self.max_pred_len = max_pred_len
        self.start_token = start_token
        self.end_token = end_token

        # encoder
        self.enc_embeddings = nn.Embedding(enc_v_dim,emb_dim)
        self.enc_embeddings.weight.data.normal_(0,0.1)
        self.encoder = nn.LSTM(emb_dim,hidden_size,1,batch_first=True)

        # decoder
        self.dec_embeddings = nn.Embedding(dec_v_dim,emb_dim)
        self.attn = nn.Linear(hidden_size,hidden_size)
        self.decoder_cell = nn.LSTMCell(emb_dim,hidden_size)
        self.decoder_dense = nn.Linear(hidden_size*2,dec_v_dim)

        self.opt = torch.optim.Adam(self.parameters(),lr=0.001)
        
    
    def encode(self,x):
        embedded = self.enc_embeddings(x)   # [batch_size, seq_len, emb_dim]
        hidden = (torch.zeros(1,x.shape[0],self.hidden_size),torch.zeros(1,x.shape[0],self.hidden_size))
        o,(h,c) = self.encoder(embedded,hidden) # [batch_size, seq_len, hidden_size], [num_layers * num_directions, batch_size, hidden_size]
        return o,h,c
    
    def inference(self,x,return_align=False):
        self.eval()
        # x [batch_size,seq_len]
        o,hx,cx = self.encode(x)    # [batch_size, seq_len, hidden_size], [num_layers * num_directions, batch_size, hidden_size] * 2
        hx,cx = hx[0],cx[0]         # [batch_size, hidden_size]
        start = torch.ones(x.shape[0],1)    # [batch_size, 1]
        start[:,0] = torch.tensor(self.start_token)
        start= start.type(torch.LongTensor)
        dec_emb_in = self.dec_embeddings(start) # [batch_size, 1, emb_dim]
        dec_emb_in = dec_emb_in.permute(1,0,2)  # [1, batch_size, emb_dim]
        dec_in = dec_emb_in[0]                  # [batch_size, emb_dim]
        output = []
        for i in range(self.max_pred_len):
            # hx.unsqueeze(1)在第一个维度上添加一个维度
            score = torch.matmul(self.attn(hx.unsqueeze(1)),o.permute(0,2,1)) # [batch_size, 1, seq_len]
            score01 = softmax(score, dim=2)  # [batch_size, 1, seq_len]
            attnDistribute = torch.matmul(score01,o)    # [batch_size, 1, hidden_size]
           
            hx, cx = self.decoder_cell(dec_in, (hx, cx))
            ha = torch.cat([attnDistribute.squeeze(1),hx],dim=1)           # [batch_size, hidden_size *2]
            result = self.decoder_dense(ha)
            result = result.argmax(dim=1).view(-1,1)
            dec_in=self.dec_embeddings(result).permute(1,0,2)[0]
            output.append(result)
        output = torch.stack(output,dim=0)
        self.train()

        return output.permute(1,0,2).view(-1,self.max_pred_len)
    
    def train_logit(self,x,y):
        o,hx,cx = self.encode(x)    # [batch_size, seq_len, hidden_size], [num_layers * num_directions, batch_size, hidden_size] * 2
        hx,cx = hx[0],cx[0]         # [batch_size, hidden_size]
        dec_in = y[:,:-1]           # [batch_size, seq_len]
        dec_emb_in = self.dec_embeddings(dec_in)    # [batch_size, seq_len, emb_dim]
        dec_emb_in = dec_emb_in.permute(1,0,2)      # [seq_len, batch_size, emb_dim]
        output = []
        for i in range(dec_emb_in.shape[0]):
            score = torch.matmul(self.attn(hx.unsqueeze(1)),o.permute(0,2,1)) # [batch_size, 1, seq_len]
            score01 = softmax(score, dim=2)  # [batch_size, 1, seq_len]
            attnDistribute = torch.matmul(score01,o)    # [batch_size, 1, hidden_size]
           
            hx, cx = self.decoder_cell(dec_emb_in[i], (hx, cx))     # [batch_size, hidden_size]
            ha = torch.cat([attnDistribute.squeeze(1),hx],dim=1)           # [batch_size, hidden_size *2]
            
            result = self.decoder_dense(ha)                              # [batch_size, dec_v_dim]
            output.append(result)
        output = torch.stack(output,dim=0)  # [seq_len, batch_size, dec_v_dim]
        return output.permute(1,0,2)        # [batch_size, seq_len, dec_v_dim]
    
    def step(self,x,y):
        self.opt.zero_grad()
        logit = self.train_logit(x,y)    
        dec_out = y[:,1:]
        loss = cross_entropy(logit.reshape(-1,self.dec_v_dim),dec_out.reshape(-1))
        loss.backward()
        self.opt.step()
        return loss.detach().numpy()

dataset = utils.DateData(4000)

loader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True
)

model = Seq2Seq(
    dataset.num_word,
    dataset.num_word,
    emb_dim=16,
    hidden_size=32,
    max_pred_len=11,
    start_token=dataset.start_token,
    end_token=dataset.end_token
)

def train():
    for i in range(100):
        for batch_idx , batch in enumerate(loader):
            bx, by, _ = batch
            loss = model.step(bx,by)
            if batch_idx % 70 == 0:
                target = dataset.idx2str(by[0, 1:-1].data.numpy())
                pred = model.inference(bx[0:1])
                res = dataset.idx2str(pred[0].data.numpy())
                src = dataset.idx2str(bx[0].data.numpy())
                print(
                    "Epoch: ",i,
                    "| t: ", batch_idx,
                    "| loss: %.3f" % loss,
                    "| input: ", src,
                    "| target: ", target,
                    "| inference: ", res,
                )
    # pkl_data = {"i2v": dataset.i2v, "x": dataset.x[:6], "y": dataset.y[:6], "align": model.inference(dataset.x[:6], return_align=True)}

    # with open("./visual/tmp/attention_align.pkl", "wb") as f:
    #     pickle.dump(pkl_data, f)

if __name__ == "__main__":
    train()