# The implementation is borrowed and modified from the official PyTorch website and ABINet: # https://pytorch.org/docs/stable/_modules/torch/nn/modules/transformer.html # https://github.com/FangShancheng/ABINet/blob/main/modules/transformer.py import copy import torch.nn as nn from torch import Tensor from torch.nn import Dropout, LayerNorm, Linear, Module, ModuleList from torch.nn import functional as F class TransformerDecoder(Module): r"""TransformerDecoder is a stack of N decoder layers Args: decoder_layer: an instance of the TransformerDecoderLayer() class (required). num_layers: the number of sub-decoder-layers in the decoder (required). norm: the layer normalization component (optional). Examples: >>> decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8) >>> transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers=6) >>> memory = torch.rand(10, 32, 512) >>> tgt = torch.rand(20, 32, 512) >>> out = transformer_decoder(tgt, memory) """ __constants__ = ['norm'] def __init__(self, decoder_layer, num_layers, norm=None): super(TransformerDecoder, self).__init__() self.layers = _get_clones(decoder_layer, num_layers) self.num_layers = num_layers self.norm = norm def forward(self, tgt, memory, memory2=None, tgt_mask=None, memory_mask=None, memory_mask2=None, tgt_key_padding_mask=None, memory_key_padding_mask=None, memory_key_padding_mask2=None): r"""Pass the inputs (and mask) through the decoder layer in turn. Args: tgt: the sequence to the decoder (required). memory: the sequence from the last layer of the encoder (required). tgt_mask: the mask for the tgt sequence (optional). memory_mask: the mask for the memory sequence (optional). tgt_key_padding_mask: the mask for the tgt keys per batch (optional). memory_key_padding_mask: the mask for the memory keys per batch (optional). Shape: see the docs in Transformer class. """ output = tgt for mod in self.layers: output = mod( output, memory, memory2=memory2, tgt_mask=tgt_mask, memory_mask=memory_mask, memory_mask2=memory_mask2, tgt_key_padding_mask=tgt_key_padding_mask, memory_key_padding_mask=memory_key_padding_mask, memory_key_padding_mask2=memory_key_padding_mask2) if self.norm is not None: output = self.norm(output) return output class TransformerDecoderLayer(Module): r"""TransformerDecoderLayer is made up of self-attn, multi-head-attn and feedforward network. Args: d_model: the number of expected features in the input (required). nhead: the number of heads in the multiheadattention models (required). dim_feedforward: the dimension of the feedforward network model (default=2048). dropout: the dropout value (default=0.1). activation: the activation function of intermediate layer, relu or gelu (default=relu). """ def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation='relu', self_attn=True, siamese=False, debug=False): super(TransformerDecoderLayer, self).__init__() self.has_self_attn, self.siamese = self_attn, siamese self.debug = debug if self.has_self_attn: self.self_attn = nn.MultiheadAttention( d_model, nhead, dropout=dropout) self.norm1 = LayerNorm(d_model) self.dropout1 = Dropout(dropout) self.multihead_attn = nn.MultiheadAttention( d_model, nhead, dropout=dropout) # Implementation of Feedforward model self.linear1 = Linear(d_model, dim_feedforward) self.dropout = Dropout(dropout) self.linear2 = Linear(dim_feedforward, d_model) self.norm2 = LayerNorm(d_model) self.norm3 = LayerNorm(d_model) self.dropout2 = Dropout(dropout) self.dropout3 = Dropout(dropout) if self.siamese: self.multihead_attn2 = nn.MultiheadAttention( d_model, nhead, dropout=dropout) self.activation = _get_activation_fn(activation) def __setstate__(self, state): if 'activation' not in state: state['activation'] = F.relu super(TransformerDecoderLayer, self).__setstate__(state) def forward(self, tgt, memory, tgt_mask=None, memory_mask=None, tgt_key_padding_mask=None, memory_key_padding_mask=None, memory2=None, memory_mask2=None, memory_key_padding_mask2=None): r"""Pass the inputs (and mask) through the decoder layer. Args: tgt: the sequence to the decoder layer (required). memory: the sequence from the last layer of the encoder (required). tgt_mask: the mask for the tgt sequence (optional). memory_mask: the mask for the memory sequence (optional). tgt_key_padding_mask: the mask for the tgt keys per batch (optional). memory_key_padding_mask: the mask for the memory keys per batch (optional). Shape: see the docs in Transformer class. """ if self.has_self_attn: tgt2, attn = self.self_attn( tgt, tgt, tgt, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask, ) tgt = tgt + self.dropout1(tgt2) tgt = self.norm1(tgt) if self.debug: self.attn = attn tgt2, attn2 = self.multihead_attn( tgt, memory, memory, attn_mask=memory_mask, key_padding_mask=memory_key_padding_mask) if self.debug: self.attn2 = attn2 if self.siamese: tgt3, attn3 = self.multihead_attn2( tgt, memory2, memory2, attn_mask=memory_mask2, key_padding_mask=memory_key_padding_mask2) tgt = tgt + self.dropout2(tgt3) if self.debug: self.attn3 = attn3 tgt = tgt + self.dropout2(tgt2) tgt = self.norm2(tgt) tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt)))) tgt = tgt + self.dropout3(tgt2) tgt = self.norm3(tgt) return tgt def _get_clones(module, N): return ModuleList([copy.deepcopy(module) for i in range(N)]) def _get_activation_fn(activation): if activation == 'relu': return F.relu elif activation == 'gelu': return F.gelu raise RuntimeError( 'activation should be relu/gelu, not {}'.format(activation))