def __init__(
self,
vocab_size: int,
encoder_output_size: int,
attention_heads: int = 4,
linear_units: int = 2048,
num_blocks: int = 6,
dropout_rate: float = 0.1,
positional_dropout_rate: float = 0.1,
self_attention_dropout_rate: float = 0.0,
src_attention_dropout_rate: float = 0.0,
input_layer: str = "embed",
use_output_layer: bool = True,
pos_enc_class=PositionalEncoding,
normalize_before: bool = True,
concat_after: bool = False,
conv_wshare: int = 4,
conv_kernel_length: Sequence[int] = (11, 11, 11, 11, 11, 11),
conv_usebias: int = False,
):
assert check_argument_types()
if len(conv_kernel_length) != num_blocks:
raise ValueError(
"conv_kernel_length must have equal number of values to num_blocks: "
f"{len(conv_kernel_length)} != {num_blocks}")
super().__init__(
vocab_size=vocab_size,
encoder_output_size=encoder_output_size,
dropout_rate=dropout_rate,
positional_dropout_rate=positional_dropout_rate,
input_layer=input_layer,
use_output_layer=use_output_layer,
pos_enc_class=pos_enc_class,
normalize_before=normalize_before,
)
attention_dim = encoder_output_size
self.decoders = repeat(
num_blocks,
lambda lnum: DecoderLayer(
attention_dim,
LightweightConvolution2D(
wshare=conv_wshare,
n_feat=attention_dim,
dropout_rate=self_attention_dropout_rate,
kernel_size=conv_kernel_length[lnum],
use_kernel_mask=True,
use_bias=conv_usebias,
),
MultiHeadedAttention(attention_heads, attention_dim,
src_attention_dropout_rate),
PositionwiseFeedForward(attention_dim, linear_units,
dropout_rate),
dropout_rate,
normalize_before,
concat_after,
),
)
def __init__(
self,
idim,
selfattention_layer_type="selfattn",
attention_dim=256,
attention_heads=4,
conv_wshare=4,
conv_kernel_length=11,
conv_usebias=False,
linear_units=2048,
num_blocks=6,
dropout_rate=0.1,
positional_dropout_rate=0.1,
attention_dropout_rate=0.0,
input_layer="conv2d",
pos_enc_class=PositionalEncoding,
normalize_before=True,
concat_after=False,
positionwise_layer_type="linear",
positionwise_conv_kernel_size=1,
padding_idx=-1,
):
"""Construct an Encoder object."""
super(Encoder, self).__init__()
self._register_load_state_dict_pre_hook(_pre_hook)
if input_layer == "linear":
self.embed = torch.nn.Sequential(
torch.nn.Linear(idim, attention_dim),
torch.nn.LayerNorm(attention_dim),
torch.nn.Dropout(dropout_rate),
torch.nn.ReLU(),
pos_enc_class(attention_dim, positional_dropout_rate),
)
elif input_layer == "conv2d":
self.embed = Conv2dSubsampling(idim, attention_dim, dropout_rate)
elif input_layer == "conv2d-scaled-pos-enc":
self.embed = Conv2dSubsampling(
idim,
attention_dim,
dropout_rate,
pos_enc_class(attention_dim, positional_dropout_rate),
)
elif input_layer == "conv2d6":
self.embed = Conv2dSubsampling6(idim, attention_dim, dropout_rate)
elif input_layer == "conv2d8":
self.embed = Conv2dSubsampling8(idim, attention_dim, dropout_rate)
elif input_layer == "vgg2l":
self.embed = VGG2L(idim, attention_dim)
elif input_layer == "embed":
self.embed = torch.nn.Sequential(
torch.nn.Embedding(idim,
attention_dim,
padding_idx=padding_idx),
pos_enc_class(attention_dim, positional_dropout_rate),
)
elif isinstance(input_layer, torch.nn.Module):
self.embed = torch.nn.Sequential(
input_layer,
pos_enc_class(attention_dim, positional_dropout_rate),
)
elif input_layer is None:
self.embed = torch.nn.Sequential(
pos_enc_class(attention_dim, positional_dropout_rate))
else:
raise ValueError("unknown input_layer: " + input_layer)
self.normalize_before = normalize_before
positionwise_layer, positionwise_layer_args = self.get_positionwise_layer(
positionwise_layer_type,
attention_dim,
linear_units,
dropout_rate,
positionwise_conv_kernel_size,
)
if selfattention_layer_type == "selfattn":
logging.info("encoder self-attention layer type = self-attention")
self.encoders = repeat(
num_blocks,
lambda lnum: EncoderLayer(
attention_dim,
MultiHeadedAttention_wordscale(
attention_heads, attention_dim, attention_dropout_rate
),
positionwise_layer(*positionwise_layer_args),
dropout_rate,
normalize_before,
concat_after,
),
)
elif selfattention_layer_type == "lightconv":
logging.info(
"encoder self-attention layer type = lightweight convolution")
self.encoders = repeat(
num_blocks,
lambda lnum: EncoderLayer(
attention_dim,
LightweightConvolution(
conv_wshare,
attention_dim,
attention_dropout_rate,
conv_kernel_length,
lnum,
use_bias=conv_usebias,
),
positionwise_layer(*positionwise_layer_args),
dropout_rate,
normalize_before,
concat_after,
),
)
elif selfattention_layer_type == "lightconv2d":
logging.info("encoder self-attention layer "
"type = lightweight convolution 2-dimentional")
self.encoders = repeat(
num_blocks,
lambda lnum: EncoderLayer(
attention_dim,
LightweightConvolution2D(
conv_wshare,
attention_dim,
attention_dropout_rate,
conv_kernel_length,
lnum,
use_bias=conv_usebias,
),
positionwise_layer(*positionwise_layer_args),
dropout_rate,
normalize_before,
concat_after,
),
)
elif selfattention_layer_type == "dynamicconv":
logging.info(
"encoder self-attention layer type = dynamic convolution")
self.encoders = repeat(
num_blocks,
lambda lnum: EncoderLayer(
attention_dim,
DynamicConvolution(
conv_wshare,
attention_dim,
attention_dropout_rate,
conv_kernel_length,
lnum,
use_bias=conv_usebias,
),
positionwise_layer(*positionwise_layer_args),
dropout_rate,
normalize_before,
concat_after,
),
)
elif selfattention_layer_type == "dynamicconv2d":
logging.info(
"encoder self-attention layer type = dynamic convolution 2-dimentional"
)
self.encoders = repeat(
num_blocks,
lambda lnum: EncoderLayer(
attention_dim,
DynamicConvolution2D(
conv_wshare,
attention_dim,
attention_dropout_rate,
conv_kernel_length,
lnum,
use_bias=conv_usebias,
),
positionwise_layer(*positionwise_layer_args),
dropout_rate,
normalize_before,
concat_after,
),
)
if self.normalize_before:
self.after_norm = LayerNorm(attention_dim)
def __init__(
self,
odim,
selfattention_layer_type="selfattn",
attention_dim=256,
attention_heads=4,
conv_wshare=4,
conv_kernel_length=11,
conv_usebias=False,
linear_units=2048,
num_blocks=6,
dropout_rate=0.1,
positional_dropout_rate=0.1,
self_attention_dropout_rate=0.0,
src_attention_dropout_rate=0.0,
input_layer="embed",
use_output_layer=True,
pos_enc_class=PositionalEncoding,
normalize_before=True,
concat_after=False,
):
"""Construct an Decoder object."""
torch.nn.Module.__init__(self)
self._register_load_state_dict_pre_hook(_pre_hook)
if input_layer == "embed":
self.embed = torch.nn.Sequential(
torch.nn.Embedding(odim, attention_dim),
pos_enc_class(attention_dim, positional_dropout_rate),
)
elif input_layer == "linear":
self.embed = torch.nn.Sequential(
torch.nn.Linear(odim, attention_dim),
torch.nn.LayerNorm(attention_dim),
torch.nn.Dropout(dropout_rate),
torch.nn.ReLU(),
pos_enc_class(attention_dim, positional_dropout_rate),
)
elif isinstance(input_layer, torch.nn.Module):
self.embed = torch.nn.Sequential(
input_layer,
pos_enc_class(attention_dim, positional_dropout_rate))
else:
raise NotImplementedError(
"only `embed` or torch.nn.Module is supported.")
self.normalize_before = normalize_before
if selfattention_layer_type == "selfattn":
logging.info("decoder self-attention layer type = self-attention")
self.decoders = repeat(
num_blocks,
lambda lnum: DecoderLayer(
attention_dim,
MultiHeadedAttention(attention_heads, attention_dim,
self_attention_dropout_rate),
MultiHeadedAttention(attention_heads, attention_dim,
src_attention_dropout_rate),
PositionwiseFeedForward(attention_dim, linear_units,
dropout_rate),
PositionwiseFeedForward(attention_dim, linear_units,
dropout_rate),
dropout_rate,
normalize_before,
concat_after,
),
)
elif selfattention_layer_type == "lightconv":
logging.info(
"decoder self-attention layer type = lightweight convolution")
self.decoders = repeat(
num_blocks,
lambda lnum: DecoderLayer(
attention_dim,
LightweightConvolution(
conv_wshare,
attention_dim,
self_attention_dropout_rate,
conv_kernel_length,
lnum,
use_kernel_mask=True,
use_bias=conv_usebias,
),
MultiHeadedAttention(attention_heads, attention_dim,
src_attention_dropout_rate),
PositionwiseFeedForward(attention_dim, linear_units,
dropout_rate),
dropout_rate,
normalize_before,
concat_after,
),
)
elif selfattention_layer_type == "lightconv2d":
logging.info("decoder self-attention layer "
"type = lightweight convolution 2-dimentional")
self.decoders = repeat(
num_blocks,
lambda lnum: DecoderLayer(
attention_dim,
LightweightConvolution2D(
conv_wshare,
attention_dim,
self_attention_dropout_rate,
conv_kernel_length,
lnum,
use_kernel_mask=True,
use_bias=conv_usebias,
),
MultiHeadedAttention(attention_heads, attention_dim,
src_attention_dropout_rate),
PositionwiseFeedForward(attention_dim, linear_units,
dropout_rate),
dropout_rate,
normalize_before,
concat_after,
),
)
elif selfattention_layer_type == "dynamicconv":
logging.info(
"decoder self-attention layer type = dynamic convolution")
self.decoders = repeat(
num_blocks,
lambda lnum: DecoderLayer(
attention_dim,
DynamicConvolution(
conv_wshare,
attention_dim,
self_attention_dropout_rate,
conv_kernel_length,
lnum,
use_kernel_mask=True,
use_bias=conv_usebias,
),
MultiHeadedAttention(attention_heads, attention_dim,
src_attention_dropout_rate),
PositionwiseFeedForward(attention_dim, linear_units,
dropout_rate),
dropout_rate,
normalize_before,
concat_after,
),
)
elif selfattention_layer_type == "dynamicconv2d":
logging.info(
"decoder self-attention layer type = dynamic convolution 2-dimentional"
)
self.decoders = repeat(
num_blocks,
lambda lnum: DecoderLayer(
attention_dim,
DynamicConvolution2D(
conv_wshare,
attention_dim,
self_attention_dropout_rate,
conv_kernel_length,
lnum,
use_kernel_mask=True,
use_bias=conv_usebias,
),
MultiHeadedAttention(attention_heads, attention_dim,
src_attention_dropout_rate),
PositionwiseFeedForward(attention_dim, linear_units,
dropout_rate),
dropout_rate,
normalize_before,
concat_after,
),
)
self.selfattention_layer_type = selfattention_layer_type
if self.normalize_before:
self.after_norm = LayerNorm(attention_dim)
if use_output_layer:
self.output_layer = torch.nn.Linear(attention_dim, odim)
else:
self.output_layer = None