(
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,
layer_drop_rate: float = 0.0,
qk_norm: bool = False,
use_flash_attn: bool = True,
gradient_checkpoint_layers: List[int] = [],
)
| 393 | class TransformerDecoder(BaseTransformerDecoder): |
| 394 | @typechecked |
| 395 | def __init__( |
| 396 | self, |
| 397 | vocab_size: int, |
| 398 | encoder_output_size: int, |
| 399 | attention_heads: int = 4, |
| 400 | linear_units: int = 2048, |
| 401 | num_blocks: int = 6, |
| 402 | dropout_rate: float = 0.1, |
| 403 | positional_dropout_rate: float = 0.1, |
| 404 | self_attention_dropout_rate: float = 0.0, |
| 405 | src_attention_dropout_rate: float = 0.0, |
| 406 | input_layer: str = "embed", |
| 407 | use_output_layer: bool = True, |
| 408 | pos_enc_class=PositionalEncoding, |
| 409 | normalize_before: bool = True, |
| 410 | concat_after: bool = False, |
| 411 | layer_drop_rate: float = 0.0, |
| 412 | qk_norm: bool = False, |
| 413 | use_flash_attn: bool = True, |
| 414 | gradient_checkpoint_layers: List[int] = [], |
| 415 | ): |
| 416 | super().__init__( |
| 417 | vocab_size=vocab_size, |
| 418 | encoder_output_size=encoder_output_size, |
| 419 | dropout_rate=dropout_rate, |
| 420 | positional_dropout_rate=positional_dropout_rate, |
| 421 | input_layer=input_layer, |
| 422 | use_output_layer=use_output_layer, |
| 423 | pos_enc_class=pos_enc_class, |
| 424 | normalize_before=normalize_before, |
| 425 | gradient_checkpoint_layers=gradient_checkpoint_layers, |
| 426 | ) |
| 427 | |
| 428 | if use_flash_attn: |
| 429 | try: |
| 430 | from espnet2.torch_utils.get_flash_attn_compatability import ( |
| 431 | is_flash_attn_supported, |
| 432 | ) |
| 433 | |
| 434 | use_flash_attn = is_flash_attn_supported() |
| 435 | import flash_attn # noqa |
| 436 | except Exception: |
| 437 | use_flash_attn = False |
| 438 | |
| 439 | attention_dim = encoder_output_size |
| 440 | self.decoders = repeat( |
| 441 | num_blocks, |
| 442 | lambda lnum: DecoderLayer( |
| 443 | attention_dim, |
| 444 | MultiHeadedAttention( |
| 445 | attention_heads, |
| 446 | attention_dim, |
| 447 | self_attention_dropout_rate, |
| 448 | qk_norm, |
| 449 | use_flash_attn, |
| 450 | True, |
| 451 | False, |
| 452 | ), |
no test coverage detected