m-e-r-c-u-r-y / pytorch-transformers Goto Github PK

import torch

class Encoder(torch.nn.Module):
    
    def __init__(self, input_dim, hid_dim, n_layers, n_heads, dropout, device, max_length):
        
        super(Encoder, self).__init__()
        
        self.device = device
        self.tok_embedding = torch.nn.Embedding(input_dim, hid_dim)
        self.pos_embedding = torch.nn.Embedding(max_length, hid_dim)
        self.layers = torch.nn.ModuleList([EncoderLayer(hid_dim, n_heads, dropout, device) for _ in range(n_layers)])
        self.dropout = torch.nn.Dropout(dropout)
        self.scale = torch.sqrt(torch.FloatTensor([hid_dim]))
        
    def forward(self, src, src_mask):
        
        # src = [batch size, src len]
        # src_mask = [batch size, src len]
        
        batch_size = src.shape[0]
        src_len = src.shape[1]
        pos = torch.arange(0, src_len).unsqueeze(0).repeat(batch_size, 1).to(self.device)
        
        # pos = [batch size, src len]
        
        src = self.dropout((self.tok_embedding(src) * self.scale.to(self.device)) + self.pos_embedding(pos))
        
        # src = [batch size, src len, hid dim]
        
        for layer in self.layers:
            src, attention = layer(src, src_mask)
        
        # src = [batch size, src len, hid dim]
        
        return src, attention

class EncoderLayer(torch.nn.Module):
    
    def __init__(self, hid_dim, n_heads, dropout, device):
        
        super(EncoderLayer, self).__init__()
        
        self.self_attn_layer_norm = torch.nn.LayerNorm(hid_dim)
        self.self_attention = MultiHeadAttentionLayer(hid_dim, n_heads, dropout, device)
        self.dropout = torch.nn.Dropout(dropout)
    
    def forward(self, src, src_mask):
        
        # src = [batch size, src len, hid dim]
        # src_mask = [batch size, src len]
                
        # self attention
        _src, attention = self.self_attention(src, src, src, src_mask)
        
        # dropout, residual connection and layer norm
        src = self.self_attn_layer_norm(src + self.dropout(_src))
        
        return src, attention

class MultiHeadAttentionLayer(torch.nn.Module):
    
    def __init__(self, hid_dim, n_heads, dropout, device):
        
        super(MultiHeadAttentionLayer, self).__init__()
        
        assert hid_dim % n_heads == 0
        
        self.device = device
        self.hid_dim = hid_dim
        self.n_heads = n_heads
        self.head_dim = hid_dim // n_heads
        self.fc_q = torch.nn.Linear(hid_dim, hid_dim)
        self.fc_k = torch.nn.Linear(hid_dim, hid_dim)
        self.fc_v = torch.nn.Linear(hid_dim, hid_dim)
        self.fc_o = torch.nn.Linear(hid_dim, hid_dim)
        self.dropout = torch.nn.Dropout(dropout)
        self.scale = torch.sqrt(torch.FloatTensor([self.head_dim]))
        
    def forward(self, query, key, value, mask = None):
        
        batch_size = query.shape[0]
        
        # query = [batch size, query len, hid dim]
        # key = [batch size, key len, hid dim]
        # value = [batch size, value len, hid dim]
        
        Q = self.fc_q(query)
        K = self.fc_k(key)
        V = self.fc_v(value)
        
        # Q = [batch size, query len, hid dim]
        # K = [batch size, key len, hid dim]
        # V = [batch size, value len, hid dim]
        
        Q = Q.view(batch_size, -1, self.n_heads, self.head_dim).permute(0, 2, 1, 3)
        K = K.view(batch_size, -1, self.n_heads, self.head_dim).permute(0, 2, 1, 3)
        V = V.view(batch_size, -1, self.n_heads, self.head_dim).permute(0, 2, 1, 3)
        
        # Q = [batch size, n heads, query len, head dim]
        # K = [batch size, n heads, key len, head dim]
        # V = [batch size, n heads, value len, head dim]
        
        energy = torch.matmul(Q, K.permute(0, 1, 3, 2)) / self.scale.to(self.device)
        
        # energy = [batch size, n heads, query len, key len]
        
        if mask is not None:
            energy.masked_fill_(mask == 0, -1e10)
        
        attention = torch.softmax(energy, dim = -1)
        
        # attention = [batch size, n heads, query len, key len]
        
        x = torch.matmul(self.dropout(attention), V)
        
        # x = [batch size, n heads, query len, head dim]
        
        x = x.permute(0, 2, 1, 3).contiguous()
        
        # x = [batch size, query len, n heads, head dim]
        
        x = x.view(batch_size, -1, self.hid_dim)
        
        # x = [batch size, query len, hid dim]
        
        x = self.fc_o(x)
        
        # x = [batch size, query len, hid dim]
        
        return x, attention

class TransformerModel(torch.nn.Module):
    
    def __init__(self, encoder, src_pad_idx, device):
        
        super(TransformerModel, self).__init__()
        
        self.encoder = encoder
        self.src_pad_idx = src_pad_idx
        self.device = device
        
    def make_src_mask(self, src):
        
        # src = [batch size, src len]
        
        src_mask = (src != self.src_pad_idx).unsqueeze(1).unsqueeze(2)

        # src_mask = [batch size, 1, 1, src len]
        
        return src_mask
    def forward(self, src):
        
        # src = [batch size, src len]
        # trg = [batch size, trg len]
        
        src_mask = self.make_src_mask(src.to(self.device))
        # src_mask = [batch size, 1, 1, src len]
        # trg_mask = [batch size, 1, trg len, trg len]
        
        enc_src, attention = self.encoder(src, src_mask)
        
        return attention



model = TransformerModel(encoder = Encoder(input_dim = 7,
                                           hid_dim = 64,
                                           n_layers = 1,
                                           n_heads = 2,
                                           dropout = 0.1,
                                           device = 'cpu',
                                           max_length = 5),
                         src_pad_idx = 0,
                         device = 'cpu').to('cpu')
a = torch.tensor([[0,2,3,5,6],[1,5,4,0,0]])
model(a)