|
|
|
|
|
|
|
|
|
|
|
|
|
"""Encoder self-attention layer definition.""" |
|
|
|
import torch |
|
|
|
from torch import nn |
|
|
|
from espnet.nets.pytorch_backend.transformer.layer_norm import LayerNorm |
|
|
|
|
|
class EncoderLayer(nn.Module): |
|
"""Encoder layer module. |
|
|
|
Args: |
|
size (int): Input dimension. |
|
self_attn (torch.nn.Module): Self-attention module instance. |
|
`MultiHeadedAttention` or `RelPositionMultiHeadedAttention` instance |
|
can be used as the argument. |
|
feed_forward (torch.nn.Module): Feed-forward module instance. |
|
`PositionwiseFeedForward`, `MultiLayeredConv1d`, or `Conv1dLinear` instance |
|
can be used as the argument. |
|
dropout_rate (float): Dropout rate. |
|
normalize_before (bool): Whether to use layer_norm before the first block. |
|
concat_after (bool): Whether to concat attention layer's input and output. |
|
if True, additional linear will be applied. |
|
i.e. x -> x + linear(concat(x, att(x))) |
|
if False, no additional linear will be applied. i.e. x -> x + att(x) |
|
|
|
""" |
|
|
|
def __init__( |
|
self, |
|
size, |
|
self_attn, |
|
feed_forward, |
|
dropout_rate, |
|
normalize_before=True, |
|
concat_after=False, |
|
): |
|
"""Construct an EncoderLayer object.""" |
|
super(EncoderLayer, self).__init__() |
|
self.self_attn = self_attn |
|
self.feed_forward = feed_forward |
|
self.norm1 = LayerNorm(size) |
|
self.norm2 = LayerNorm(size) |
|
self.dropout = nn.Dropout(dropout_rate) |
|
self.size = size |
|
self.normalize_before = normalize_before |
|
self.concat_after = concat_after |
|
if self.concat_after: |
|
self.concat_linear = nn.Linear(size + size, size) |
|
|
|
def forward(self, x, mask, cache=None): |
|
"""Compute encoded features. |
|
|
|
Args: |
|
x_input (torch.Tensor): Input tensor (#batch, time, size). |
|
mask (torch.Tensor): Mask tensor for the input (#batch, time). |
|
cache (torch.Tensor): Cache tensor of the input (#batch, time - 1, size). |
|
|
|
Returns: |
|
torch.Tensor: Output tensor (#batch, time, size). |
|
torch.Tensor: Mask tensor (#batch, time). |
|
|
|
""" |
|
residual = x |
|
if self.normalize_before: |
|
x = self.norm1(x) |
|
|
|
if cache is None: |
|
x_q = x |
|
else: |
|
assert cache.shape == (x.shape[0], x.shape[1] - 1, self.size) |
|
x_q = x[:, -1:, :] |
|
residual = residual[:, -1:, :] |
|
mask = None if mask is None else mask[:, -1:, :] |
|
|
|
if self.concat_after: |
|
x_concat = torch.cat((x, self.self_attn(x_q, x, x, mask)), dim=-1) |
|
x = residual + self.concat_linear(x_concat) |
|
else: |
|
x = residual + self.dropout(self.self_attn(x_q, x, x, mask)) |
|
if not self.normalize_before: |
|
x = self.norm1(x) |
|
|
|
residual = x |
|
if self.normalize_before: |
|
x = self.norm2(x) |
|
x = residual + self.dropout(self.feed_forward(x)) |
|
if not self.normalize_before: |
|
x = self.norm2(x) |
|
|
|
if cache is not None: |
|
x = torch.cat([cache, x], dim=1) |
|
|
|
return x, mask |
|
|
