def __init__(self, args, no_encoder_attn=False): super().__init__() self.embed_dim = args.decoder_embed_dim self.dropout = args.dropout self.relu_dropout = args.relu_dropout self.more_dropouts = args.decoder_aan_more_dropouts if args.decoder_attn_window_size <= 0: self.avg_attn = AverageAttention(self.embed_dim, dropout=args.attention_dropout) else: self.avg_attn = AverageWindowAttention( self.embed_dim, dropout=args.attention_dropout, window_size=args.decoder_attn_window_size, ) # self.activation = getattr(args, "decoder_ffn_activation", "relu") self.aan_layer_norm = LayerNorm(self.embed_dim) if args.no_decoder_aan_ffn: self.aan_ffn = None else: aan_ffn_hidden_dim = (args.decoder_ffn_embed_dim if args.decoder_aan_ffn_use_embed_dim else args.decoder_ffn_embed_dim) self.aan_ffn = FeedForwardNetwork( self.embed_dim, aan_ffn_hidden_dim, self.embed_dim, num_layers=2, dropout=args.relu_dropout, ) if args.no_decoder_aan_gating: self.aan_gating_fc = None else: self.aan_gating_fc = Linear(self.embed_dim * 2, self.embed_dim * 2) self.normalize_before = args.decoder_normalize_before if no_encoder_attn: self.encoder_attn = None self.encoder_attn_layer_norm = None else: self.encoder_attn = MultiheadAttention( self.embed_dim, args.decoder_attention_heads, kdim=args.encoder_embed_dim, vdim=args.encoder_embed_dim, dropout=args.attention_dropout, encoder_decoder_attention=True, ) self.encoder_attn_layer_norm = LayerNorm(self.embed_dim) self.ffn = FeedForwardNetwork( self.embed_dim, args.decoder_ffn_embed_dim, self.embed_dim, num_layers=2, dropout=args.relu_dropout, ) self.final_layer_norm = LayerNorm(self.embed_dim) self.need_attn = True self.onnx_trace = False
class TransformerAANDecoderLayer(nn.Module): """Decoder layer block. In the original paper each operation (multi-head attention, encoder attention or FFN) is postprocessed with: `dropout -> add residual -> layernorm`. In the tensor2tensor code they suggest that learning is more robust when preprocessing each layer with layernorm and postprocessing with: `dropout -> add residual`. We default to the approach in the paper, but the tensor2tensor approach can be enabled by setting *args.decoder_normalize_before* to ``True``. Args: args (argparse.Namespace): parsed command-line arguments no_encoder_attn (bool, optional): whether to attend to encoder outputs. Default: ``False`` """ def __init__(self, args, no_encoder_attn=False): super().__init__() self.embed_dim = args.decoder_embed_dim self.dropout = args.dropout self.relu_dropout = args.relu_dropout self.more_dropouts = args.decoder_aan_more_dropouts if args.decoder_attn_window_size <= 0: self.avg_attn = AverageAttention(self.embed_dim, dropout=args.attention_dropout) else: self.avg_attn = AverageWindowAttention( self.embed_dim, dropout=args.attention_dropout, window_size=args.decoder_attn_window_size, ) # self.activation = getattr(args, "decoder_ffn_activation", "relu") self.aan_layer_norm = LayerNorm(self.embed_dim) if args.no_decoder_aan_ffn: self.aan_ffn = None else: aan_ffn_hidden_dim = (args.decoder_ffn_embed_dim if args.decoder_aan_ffn_use_embed_dim else args.decoder_ffn_embed_dim) self.aan_ffn = FeedForwardNetwork( self.embed_dim, aan_ffn_hidden_dim, self.embed_dim, num_layers=2, dropout=args.relu_dropout, ) if args.no_decoder_aan_gating: self.aan_gating_fc = None else: self.aan_gating_fc = Linear(self.embed_dim * 2, self.embed_dim * 2) self.normalize_before = args.decoder_normalize_before if no_encoder_attn: self.encoder_attn = None self.encoder_attn_layer_norm = None else: self.encoder_attn = MultiheadAttention( self.embed_dim, args.decoder_attention_heads, kdim=args.encoder_embed_dim, vdim=args.encoder_embed_dim, dropout=args.attention_dropout, encoder_decoder_attention=True, ) self.encoder_attn_layer_norm = LayerNorm(self.embed_dim) self.ffn = FeedForwardNetwork( self.embed_dim, args.decoder_ffn_embed_dim, self.embed_dim, num_layers=2, dropout=args.relu_dropout, ) self.final_layer_norm = LayerNorm(self.embed_dim) self.need_attn = True self.onnx_trace = False def prepare_for_onnx_export_(self): self.onnx_trace = True def forward( self, x, encoder_out, encoder_padding_mask, incremental_state, prev_self_attn_state=None, prev_attn_state=None, self_attn_mask=None, self_attn_padding_mask=None, ): """ Args: x (Tensor): input to the layer of shape `(seq_len, batch, embed_dim)` encoder_padding_mask (ByteTensor): binary ByteTensor of shape `(batch, src_len)` where padding elements are indicated by ``1``. Returns: encoded output of shape `(batch, src_len, embed_dim)` """ residual = x if "residual" in self.more_dropouts: residual = F.dropout(residual, p=self.dropout, training=self.training) if prev_self_attn_state is not None: if incremental_state is None: incremental_state = {} prev_key, prev_value = prev_self_attn_state saved_state = {"prev_key": prev_key, "prev_value": prev_value} self.avg_attn._set_input_buffer(incremental_state, saved_state) x, _ = self.avg_attn( value=x, mask_future_timesteps=True, incremental_state=incremental_state, mask_trick=self.training, ) if "after_avg" in self.more_dropouts: x = F.dropout(x, p=self.dropout, training=self.training) if self.aan_layer_norm is not None: x = self.maybe_layer_norm(self.aan_layer_norm, x, before=True) if self.aan_ffn is not None: x = self.aan_ffn(x) if "after_ffn" in self.more_dropouts: x = F.dropout(x, p=self.dropout, training=self.training) if self.aan_gating_fc is not None: i, f = self.aan_gating_fc(torch.cat([residual, x], dim=-1)).chunk(2, dim=-1) x = torch.sigmoid(f) * residual + torch.sigmoid(i) * x if "after_gating" in self.more_dropouts: x = F.dropout(x, p=self.dropout, training=self.training) x = F.dropout(x, p=self.dropout, training=self.training) x = residual + x if self.aan_layer_norm is not None: x = self.maybe_layer_norm(self.aan_layer_norm, x, after=True) attn = None if self.encoder_attn is not None: residual = x x = self.maybe_layer_norm(self.encoder_attn_layer_norm, x, before=True) if prev_attn_state is not None: if incremental_state is None: incremental_state = {} prev_key, prev_value = prev_attn_state saved_state = {"prev_key": prev_key, "prev_value": prev_value} self.encoder_attn._set_input_buffer(incremental_state, saved_state) x, attn = self.encoder_attn( query=x, key=encoder_out, value=encoder_out, key_padding_mask=encoder_padding_mask, incremental_state=incremental_state, static_kv=True, need_weights=(not self.training and self.need_attn), ) x = F.dropout(x, p=self.dropout, training=self.training) x = residual + x x = self.maybe_layer_norm(self.encoder_attn_layer_norm, x, after=True) residual = x x = self.maybe_layer_norm(self.final_layer_norm, x, before=True) x = self.ffn(x) x = F.dropout(x, p=self.dropout, training=self.training) x = residual + x x = self.maybe_layer_norm(self.final_layer_norm, x, after=True) return x, attn def maybe_layer_norm(self, layer_norm, x, before=False, after=False): assert before ^ after if after ^ self.normalize_before: return layer_norm(x) else: return x def make_generation_fast_(self, need_attn=False, **kwargs): self.need_attn = need_attn def extra_repr(self): return "dropout={}, more_dropouts={}".format(self.dropout, self.more_dropouts)