文章目录
1. 详解
2. 代码
import torch
from torch import nn
from torch.nn import functional as F
class Seq2SeqEncoder(nn.Module):
def __init__(self, embedding_dim, hidden_size, source_vocab_size):
super(Seq2SeqEncoder, self).__init__()
self.lstm_layer = nn.LSTM(input_size=embedding_dim,
hidden_size=hidden_size,
batch_first=True)
self.embedding_table = nn.Embedding(source_vocab_size, embedding_dim)
def forward(self, input_idx):
# input_idx：表示原序列的句子长度
# input_sequence.shape = [bs,source_length,embedding_dim]
input_sequence = self.embedding_table(input_idx)
output_state, (final_h, final_c) = self.lstm_layer(input_sequence)
# output_state表示的是每个向量的上下文特征
return output_state, final_h
class Seq2SeqAttentionMechanism(nn.Module):
def __init__(self):
super(Seq2SeqAttentionMechanism, self).__init__()
# 解码器的第t时刻表示的状态需要跟编码器的每个时刻进行表示
def forward(self, decoder_state_t, encoder_states):
bs, source_length, hidden_size = encoder_states.shape
decoder_state_t = decoder_state_t.unsqueeze(1)
decoder_state_t = torch.tile(decoder_state_t, dims=(1, source_length, 1))
# score =[bs,source_length]
score = torch.sum(decoder_state_t * encoder_states, dim=-1)
# attn_prob = [bs,source_length]
attn_prob = F.softmax(score, dim=-1)
# context=[bs,hidden_size];第t时刻的解码器上下文向量
context = torch.sum(attn_prob.unsqueeze(-1) * encoder_states, 1)
return attn_prob, context
class Seq2SeqDecoder(nn.Module):
def __init__(self, embedding_dim, hidden_size, num_classes, target_vocab_size, start_id, end_id):
super(Seq2SeqDecoder, self).__init__()
self.lstm_cell = torch.nn.LSTMCell(embedding_dim, hidden_size)
self.proj_layer = nn.Linear(hidden_size * 2, num_classes)
self.attention_mechanism = Seq2SeqAttentionMechanism()
self.num_classes = num_classes
self.embedding_table = torch.nn.Embedding(target_vocab_size, embedding_dim)
self.start_id = start_id
self.end_id = end_id
def forward(self, shifted_target_ids, encoder_states):
# 训练阶段使用,teacher-force mode
shifted_target = self.embedding_table(shifted_target_ids)
bs, target_length, embedding_dim = shifted_target.shape
bs, source_length, hidden_size = encoder_states.shape
logits = torch.zeros(bs, target_length, self.num_classes)
probs = torch.zeros(bs, target_length, source_length)
for t in range(target_length):
# decoder=[bs,embedding_dim]
decoder_input_t = shifted_target[:, t, :]
if t == 0:
h_t, c_t = self.lstm_cell(decoder_input_t)
else:
h_t, c_t = self.lstm_cell(decoder_input_t, (h_t, c_t))
attn_prob, context = self.attention_mechanism(h_t, encoder_states)
decoder_output = torch.cat((context, h_t), -1)
logits[:, t, :] = self.proj_layer(decoder_output)
probs[:, t, :] = attn_prob
return probs, logits
def inference(self, encoder_state):
# 推理阶段使用
target_id = self.start_id
h_t = None
result = []
while True:
decoder_input = self.embedding_table(target_id)
if h_t is None:
h_t, c_t = self.lstm_cell(decoder_input)
else:
h_t, c_t = self.lstm_cell(decoder_input, (h_t, c_t))
attn_prob, context = self.attention_mechanism(h_t, encoder_state)
decoder_output = torch.cat((context, h_t), -1)
logits = self.proj_layer(decoder_output)
target_id = torch.argmax(logits, -1)
result.append(target_id)
if torch.any(target_id == self.end_id):
print("stop decoding")
break
predicted_ids = torch.stack(result, dim=0)
return predicted_ids
class Model(nn.Module):
def __init__(self, embedding_dim, hidden_size, num_classes,
source_vocab_size, target_vocab_size, start_id, end_id):
super(Model, self).__init__()
self.encoder = Seq2SeqEncoder(embedding_dim, hidden_size, source_vocab_size)
self.decoder = Seq2SeqDecoder(embedding_dim, hidden_size, num_classes,
target_vocab_size, start_id, end_id)
def forward(self, input_sequence_ids, shifted_target_ids):
# 训练阶段
encoder_states, final_h = self.encoder(input_sequence_ids)
prob, logits = self.decoder(shifted_target_ids, encoder_states)
return prob, logits
def infer(self):
pass
if __name__ == "__main__":
source_length = 3
target_length = 4
embedding_dim = 8
hidden_size = 16
num_classes = 10
bs = 2
start_id = end_id = 0
source_vocab_size = 100
target_vocab_size = 100
input_sequence_ids = torch.randint(source_vocab_size, size=(bs, source_length)).to(torch.int32)
target_ids = torch.randint(target_vocab_size, size=(bs, target_length))
target_ids = torch.cat((target_ids, end_id * torch.ones(bs, 1), target_ids[:, 1:]), dim=1).to(torch.int32)
shifted_target_ids = torch.cat((start_id * torch.ones(bs, 1), target_ids[:, 1:]), dim=1).to(torch.int32)
model = Model(embedding_dim, hidden_size, num_classes, source_vocab_size, target_vocab_size, start_id, end_id)
probs, logits = model(input_sequence_ids, shifted_target_ids)
print(f"probs={probs}")
print(f"logits={logits}")-
