wehub-resource-sync
73 lines
2.9 KiB
Python
73 lines
2.9 KiB
Python
# Copyright (c) Facebook, Inc. and its affiliates.
|
|
#
|
|
# This source code is licensed under the MIT license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
import torch.nn as nn
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
|
|
class LocationAttention(nn.Module):
|
|
"""
|
|
Attention-Based Models for Speech Recognition
|
|
https://arxiv.org/pdf/1506.07503.pdf
|
|
|
|
:param int encoder_dim: # projection-units of encoder
|
|
:param int decoder_dim: # units of decoder
|
|
:param int attn_dim: attention dimension
|
|
:param int conv_dim: # channels of attention convolution
|
|
:param int conv_kernel_size: filter size of attention convolution
|
|
"""
|
|
|
|
def __init__(self, attn_dim, encoder_dim, decoder_dim,
|
|
attn_state_kernel_size, conv_dim, conv_kernel_size,
|
|
scaling=2.0):
|
|
super(LocationAttention, self).__init__()
|
|
self.attn_dim = attn_dim
|
|
self.decoder_dim = decoder_dim
|
|
self.scaling = scaling
|
|
self.proj_enc = nn.Linear(encoder_dim, attn_dim)
|
|
self.proj_dec = nn.Linear(decoder_dim, attn_dim, bias=False)
|
|
self.proj_attn = nn.Linear(conv_dim, attn_dim, bias=False)
|
|
self.conv = nn.Conv1d(attn_state_kernel_size, conv_dim,
|
|
2 * conv_kernel_size + 1,
|
|
padding=conv_kernel_size, bias=False)
|
|
self.proj_out = nn.Sequential(nn.Tanh(), nn.Linear(attn_dim, 1))
|
|
|
|
self.proj_enc_out = None # cache
|
|
|
|
def clear_cache(self):
|
|
self.proj_enc_out = None
|
|
|
|
def forward(self, encoder_out, encoder_padding_mask, decoder_h, attn_state):
|
|
"""
|
|
:param torch.Tensor encoder_out: padded encoder hidden state B x T x D
|
|
:param torch.Tensor encoder_padding_mask: encoder padding mask
|
|
:param torch.Tensor decoder_h: decoder hidden state B x D
|
|
:param torch.Tensor attn_prev: previous attention weight B x K x T
|
|
:return: attention weighted encoder state (B, D)
|
|
:rtype: torch.Tensor
|
|
:return: previous attention weights (B x T)
|
|
:rtype: torch.Tensor
|
|
"""
|
|
bsz, seq_len, _ = encoder_out.size()
|
|
if self.proj_enc_out is None:
|
|
self.proj_enc_out = self.proj_enc(encoder_out)
|
|
|
|
# B x K x T -> B x C x T
|
|
attn = self.conv(attn_state)
|
|
# B x C x T -> B x T x C -> B x T x D
|
|
attn = self.proj_attn(attn.transpose(1, 2))
|
|
|
|
if decoder_h is None:
|
|
decoder_h = encoder_out.new_zeros(bsz, self.decoder_dim)
|
|
dec_h = self.proj_dec(decoder_h).view(bsz, 1, self.attn_dim)
|
|
|
|
out = self.proj_out(attn + self.proj_enc_out + dec_h).squeeze(2)
|
|
out.masked_fill_(encoder_padding_mask, -float("inf"))
|
|
|
|
w = F.softmax(self.scaling * out, dim=1)
|
|
c = torch.sum(encoder_out * w.view(bsz, seq_len, 1), dim=1)
|
|
return c, w
|