赞
踩
class Encoder(nn.Module):
def __init__(self, vocab_size, embed_size, num_hiddens, num_layers,
drop_prob=0, **kwargs):
super(Encoder, self).__init__(**kwargs)
self.embedding = nn.Embedding(vocab_size, embed_size)
self.rnn = nn.GRU(embed_size, num_hiddens, num_layers, dropout=drop_prob)
def forward(self, inputs, state):
# 输入形状是(批量大小, 时间步数)。将输出互换样本维和时间步维
embedding = self.embedding(inputs.long()).permute(1, 0, 2) # (seq_len, batch, input_size)
return self.rnn(embedding, state)
def begin_state(self):
return None # 隐藏态初始化为None时PyTorch会自动初始化为0
首先输入形状[batch,max_len]很容易理解,然后转化为词向量并交换维度得到[max_len,batch,embedd_size],经过GRU之后得到的outputs形状是[max_len,batch,num_hiddens],state形状是[num_layers,batch,num_hiddens]是最后一步得到的隐藏状态。
encoder = Encoder(vocab_size=10, embed_size=8, num_hiddens=16, num_layers=2)
output, state = encoder(torch.zeros((4, 7)), encoder.begin_state())
output.shape, state.shape # GRU的state是h, 而LSTM的是一个元组(h, c)
输出
(torch.Size([7, 4, 16]), torch.Size([2, 4, 16]))
def attention_model(input_size, attention_size):
model = nn.Sequential(nn.Linear(input_size,
attention_size, bias=False),
nn.Tanh(),
nn.Linear(attention_size, 1, bias=False))
return model
我们当然也可以采用内积,这里实现的是第二种a函数定义方式。
# 根据注意力机制返回背景变量[批量大小, 隐藏单元个数]
def attention_forward(model, enc_states, dec_state):
"""
enc_states: (时间步数, 批量大小, 隐藏单元个数) 是编码器在所有时间步的隐藏状态
dec_state: (批量大小, 隐藏单元个数) 是解码器在上一步的隐藏状态
"""
# 将解码器隐藏状态广播到和编码器隐藏状态形状相同后进行连结
dec_states = dec_state.unsqueeze(dim=0).expand_as(enc_states)# (时间步数, 批量大小, 隐藏单元个数)
enc_and_dec_states = torch.cat((enc_states, dec_states), dim=2)# (时间步数, 批量大小, 2*隐藏单元个数)
e = model(enc_and_dec_states) # 形状为(时间步数, 批量大小, 1)
alpha = F.softmax(e, dim=0) # 在时间步维度做softmax运算 形状为(时间步数, 批量大小, 1)
print((alpha * enc_states).shape) # torch.Size([7, 2, 64])
return (alpha * enc_states).sum(dim=0) # 返回背景变量 [batch,num_hiddens]
这里的dec_state就是咱们上面讲的
s
t
′
−
1
s_{t'-1}
st′−1,形状是[batch,num_hiddens],enc_states就是咱们上面讲的编码器的隐藏状态
H
H
H,形状是[max_len,batch,num_hiddens],dec_state.unsqueeze(dim=0)变为[1,batch,num_hiddens],expand_as(enc_states)变为[max_len,batch,num_hiddens],是通过复制的形式变成这个尺寸。然后在第三维度上联结,这样的话相当于把输入s和H一起放到函数a里。这里的a函数可不是计算内积了,而是第二种方法。执行完a函数得到e之后再对其的时间步维度执行softmax,然后将softmax之后的概率与隐藏层状态enc_states相乘得到每个时间步的加权隐藏层状态,再对时间步维度求和就可以得到最终的背景变量c,形状为[batch,num_hiddens]。
下面示范一下:
在下面的例子中,编码器的时间步数为10,批量大小为4,编码器和解码器的隐藏单元个数均为8。注意力机制返回一个小批量的背景向量,每个背景向量的长度等于编码器的隐藏单元个数。因此输出的形状为(4, 8)。
seq_len, batch_size, num_hiddens = 10, 4, 8
model = attention_model(2*num_hiddens, 10)
enc_states = torch.zeros((seq_len, batch_size, num_hiddens))
dec_state = torch.zeros((batch_size, num_hiddens))
attention_forward(model, enc_states, dec_state).shape # torch.Size([4, 8])
# 解码器首先计算背景向量,然后和编码后的输入联结得到(批量大小, num_hiddens+embed_size),然后放入GRU+Linear得到[batch,vocab_size] # 即从输入[batch,]得到输出[batch,vocab_size] class Decoder(nn.Module): def __init__(self, vocab_size, embed_size, num_hiddens, num_layers, attention_size, drop_prob=0): super(Decoder, self).__init__() self.embedding = nn.Embedding(vocab_size, embed_size) self.attention = attention_model(2*num_hiddens, attention_size) # GRU的输入包含attention输出的c和实际输入, 所以尺寸是 num_hiddens+embed_size self.rnn = nn.GRU(num_hiddens + embed_size, num_hiddens, num_layers, dropout=drop_prob) self.out = nn.Linear(num_hiddens, vocab_size) def forward(self, cur_input, state, enc_states): """ cur_input shape: (batch, ) state shape: (num_layers, batch, num_hiddens) """ # 使用注意力机制计算背景向量 [batch,num_hiddens] c = attention_forward(self.attention, enc_states, state[-1]) # 将嵌入后的输入和背景向量在特征维连结, (批量大小, num_hiddens+embed_size) input_and_c = torch.cat((self.embedding(cur_input), c), dim=1) # 为输入和背景向量的连结增加时间步维,时间步个数为1 rnn([(1,批量大小, num_hiddens+embed_size)],(num_layers, batch, num_hiddens)) # =(1,批量大小,num_hiddens),(num_layers, batch, num_hiddens) output, state = self.rnn(input_and_c.unsqueeze(0), state) # 移除时间步维,输出形状为(批量大小, 输出词典大小) output = self.out(output).squeeze(dim=0) return output, state def begin_state(self, enc_state): # 直接将编码器最终时间步的隐藏状态作为解码器的初始隐藏状态 return enc_state
为啥输入cur_input是(batch, )形状呢,因为每一次时间步调用一次解码器,所以输入的是某个时间步的batch个字,首先要得到背景向量,enc_states, state[-1]分别是编码器隐藏状态和解码器上一步的隐藏状态。然后还记得我们说过GRU需要改造一下,就是将c和input联结一下,然后作为输入输入到GRU层,input_and_c.unsqueeze(0)是增加一个步数维。
def batch_loss(encoder, decoder, X, Y, loss): batch_size = X.shape[0] enc_state = encoder.begin_state() # 获取特征enc_outputs,[步数,batch,num_hiddens],最后一步的隐藏层状态enc_state[num_layers,batch,num_hiddens] enc_outputs, enc_state = encoder(X, enc_state) # 初始化解码器的隐藏状态 dec_state = decoder.begin_state(enc_state) # 解码器在最初时间步的输入是BOS,输入形状为[batch,] dec_input = torch.tensor([out_vocab.stoi[BOS]] * batch_size) # 我们将使用掩码变量mask来忽略掉标签为填充项PAD的损失 mask, num_not_pad_tokens = torch.ones(batch_size,), 0 l = torch.tensor([0.0]) for y in Y.permute(1,0): # Y shape: (batch, seq_len) y:[1,batch],一个字一个字来 # 得到第一步batch个字的输出,为[batch,vocab_size],state:(num_layers, batch, num_hiddens) dec_output, dec_state = decoder(dec_input, dec_state, enc_outputs) l = l + (mask * loss(dec_output, y)).sum() dec_input = y # 使用强制教学 num_not_pad_tokens += mask.sum().item() # EOS后面全是PAD. 下面一行保证一旦遇到EOS接下来的循环中mask就一直是0 mask = mask * (y != out_vocab.stoi[EOS]).float() return l / num_not_pad_tokens def train(encoder, decoder, dataset, lr, batch_size, num_epochs): enc_optimizer = torch.optim.Adam(encoder.parameters(), lr=lr) dec_optimizer = torch.optim.Adam(decoder.parameters(), lr=lr) loss = nn.CrossEntropyLoss(reduction='none') data_iter = Data.DataLoader(dataset, batch_size, shuffle=True) for epoch in range(num_epochs): l_sum = 0.0 for X, Y in data_iter:# (batch, seq_len) enc_optimizer.zero_grad() dec_optimizer.zero_grad() l = batch_loss(encoder, decoder, X, Y, loss) l.backward() enc_optimizer.step() dec_optimizer.step() l_sum += l.item() if (epoch + 1) % 10 == 0: print("epoch %d, loss %.3f" % (epoch + 1, l_sum / len(data_iter))) embed_size, num_hiddens, num_layers = 64, 64, 2 attention_size, drop_prob, lr, batch_size, num_epochs = 10, 0.5, 0.01, 2, 50 encoder = Encoder(len(in_vocab), embed_size, num_hiddens, num_layers, drop_prob) decoder = Decoder(len(out_vocab), embed_size, num_hiddens, num_layers, attention_size, drop_prob) train(encoder, decoder, dataset, lr, batch_size, num_epochs)
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。