transformer.py

# https://github.com/guocheng2018/Transformer-Encoder/tree/master/transformer_encoder

import copy
import math
import torch
import torch.nn as nn
import torch.nn.functional as F

def clones(module, N):
    "Produce N identical layers."
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

class ScaledDotProductAttention(nn.Module):
    def __init__(self):
        super(ScaledDotProductAttention, self).__init__()

    def forward(self, query, key, value, mask = None, dropout = None):
        """
        Args:
            `query`: shape (batch_size, n_heads, max_len, d_q)
            `key`: shape (batch_size, n_heads, max_len, d_k)
            `value`: shape (batch_size, n_heads, max_len, d_v)
            `mask`: shape (batch_size, 1, 1, max_len)
            `dropout`: nn.Dropout
        Returns:
            `weighted value`: shape (batch_size, n_heads, max_len, d_v)
            `weight matrix`: shape (batch_size, n_heads, max_len, max_len)
        """
        d_k         = query.size(-1)  # d_k = d_model / n_heads
        scores      = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)  # B*H*L*L
        
        if mask is not None:
            scores = scores.masked_fill(mask.eq(0), -1e9)
        
        p_attn      = F.softmax(scores, dim = -1)  # B*H*L*L
        
        if dropout is not None:
            p_attn = dropout(p_attn)

        return torch.matmul(p_attn, value), p_attn

class MultiHeadAttention(nn.Module):

    def __init__(self, n_heads, d_model, dropout = 0.1):
        super(MultiHeadAttention, self).__init__()

        assert d_model % n_heads == 0
        # We assume d_v always equals d_k
        self.d_k        = d_model // n_heads
        self.h          = n_heads
        self.linears    = clones(nn.Linear(d_model, d_model), 4)
        self.sdpa       = ScaledDotProductAttention()
        self.attn       = None
        self.dropout    = nn.Dropout(p = dropout)

    def forward(self, query, key, value, mask = None):
        """
        Args: 
            `query`: shape (batch_size, max_len, d_model)
            `key`: shape (batch_size, max_len, d_model)
            `value`: shape (batch_size, max_len, d_model)
            `mask`: shape (batch_size, max_len)
        
        Returns:
            shape (batch_size, max_len, d_model)
        """
        if mask is not None:
            # Same mask applied to all h heads. B*1*1*L
            mask = mask.unsqueeze(1).unsqueeze(1)
        batch_size = query.size(0)

        # 1) Do all the linear projections in batch from d_model => h x d_k
        query, key, value = [l(x).view(batch_size, -1, self.h, self.d_k).transpose(1, 2) for l, x in zip(self.linears, (query, key, value))]

        # 2) Apply attention on all the projected vectors in batch.
        # x: B x H x L x D_v
        x, self.attn = self.sdpa(query, key, value, mask = mask, dropout = self.dropout)

        # 3) "Concat" using a view and apply a final linear.
        x = x.transpose(1, 2).contiguous().view(batch_size, -1, self.h * self.d_k)

        return self.linears[-1](x)

class FeedForward(nn.Module):

    def __init__(self, d_model, d_ff, dropout = 0.1):
        super(FeedForward, self).__init__()

        self.w_1        = nn.Linear(d_model, d_ff)
        self.w_2        = nn.Linear(d_ff, d_model)
        self.dropout    = nn.Dropout(dropout)

    def forward(self, x):
        """
        Args:
            `x`: shape (batch_size, max_len, d_model)
        Returns:
            same shape as input x
        """
        return self.w_2(self.dropout(F.relu(self.w_1(x))))

class LayerNorm(nn.Module):

    def __init__(self, features, eps = 1e-6):
        # features = d_model
        super(LayerNorm, self).__init__()

        self.a      = nn.Parameter(torch.ones(features))
        self.b      = nn.Parameter(torch.zeros(features))
        self.eps    = eps

    def forward(self, x):
        mean    = x.mean(-1, keepdim=True)
        std     = x.std(-1, keepdim=True)

        return self.a * (x - mean) / (std + self.eps) + self.b

class SublayerConnection(nn.Module):
    """
    A residual connection followed by a layer norm.
    Note for code simplicity the norm is first as opposed to last.
    """

    def __init__(self, size, dropout):
        super(SublayerConnection, self).__init__()

        self.norm       = LayerNorm(size)
        self.dropout    = nn.Dropout(dropout)

    def forward(self, x, sublayer):
        """Apply residual connection to any sublayer with the same size."""
        return x + self.dropout(sublayer(self.norm(x)))

class EncoderLayer(nn.Module):
    """Encoder is made up of self-attn and feed forward"""

    def __init__(self, size, self_attn, feed_forward, dropout):
        super(EncoderLayer, self).__init__()

        self.self_attn      = self_attn
        self.feed_forward   = feed_forward
        self.sublayer       = clones(SublayerConnection(size, dropout), 2)
        self.size           = size

    def forward(self, x, mask):
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))

        return self.sublayer[1](x, self.feed_forward)

class Encoder(nn.Module):
    "Core encoder is a stack of N layers"

    def __init__(self, layer, N):
        super(Encoder, self).__init__()

        self.layers = clones(layer, N)
        self.norm 	= LayerNorm(layer.size)

    def forward(self, x, mask):
        "Pass the input (and mask) through each layer in turn."
        for layer in self.layers:
            x = layer(x, mask)
        return self.norm(x)

class TransformerEncoder(nn.Module):
    """The encoder of transformer
    Args:
        `n_layers`: number of stacked encoder layers
        `d_model`: model dimension
        `d_ff`: hidden dimension of feed forward layer
        `n_heads`: number of heads of self-attention
        `dropout`: dropout rate, default 0.1
    """

    def __init__(self, d_model, d_ff, n_heads, n_layers, dropout = 0.1):
        super(TransformerEncoder, self).__init__()

        self.multi_headed_attention = MultiHeadAttention(n_heads, d_model, dropout)
        self.feed_forward 			= FeedForward(d_model, d_ff, dropout)
        self.encoder_layer 			= EncoderLayer(d_model, self.multi_headed_attention, self.feed_forward, dropout)
        self.encoder 				= Encoder(self.encoder_layer, n_layers)
        self.reset_parameters()

    def reset_parameters(self):
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)

    def forward(self, x, mask):
        return self.encoder(x, mask)

# class PositionalEncoding(nn.Module):

#     def __init__(self, d_model, dropout = 0.1, max_len = 12):
#         super(PositionalEncoding, self).__init__()

#         self.dropout = nn.Dropout(p = dropout)

#         # Compute the positional encodings once in log space.
#         pe          = torch.zeros(max_len, d_model)
#         position    = torch.arange(0, max_len).unsqueeze(1).float()
#         div_term    = torch.exp(torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model))
#         pe[:, 0::2] = torch.sin(position * div_term)
#         pe[:, 1::2] = torch.cos(position * div_term)
#         pe          = pe.unsqueeze(0)
#         self.register_buffer("pe", pe)

#     def forward(self, x):
#         """
#         Args:
#             x: `embeddings`, shape (batch, max_len, d_model)
#         Returns:
#             `encoder input`, shape (batch, max_len, d_model)
#         """
#         x = x + self.pe[:, : x.size(1)]

#         return self.dropout(x)

class PositionalEncoding(nn.Module):

    def __init__(self, d_model, dropout = 0.1):
        super(PositionalEncoding, self).__init__()

        self.d_model = d_model
        self.dropout = nn.Dropout(p = dropout)

    def forward(self, x, pos):
        """
        Args:
            x: `embeddings`, shape (batch, len, d_model)
        Returns:
            `encoder input`, shape (batch, len, d_model)
        """
        # B, l(11 + 1), d_model
        _, _, d_model = x.shape

        pe              = torch.zeros_like(x)
        pos             = pos.unsqueeze(-1)
        div_term        = torch.exp(torch.arange(0, d_model, 2, device = x.device).float() * -(math.log(10000.0) / d_model))
        div_term        = div_term.unsqueeze(0).unsqueeze(0)
        pe[:, :, 0::2]  = torch.sin(pos * div_term)
        pe[:, :, 1::2]  = torch.cos(pos * div_term)

        x = x + pe

        return self.dropout(x)