transformer的pytorch实现（可以根据自己的需求进行一些小trick）

客官°小女子只卖身不卖艺 2022-12-29 06:50 9阅读 0赞

根据自己对transformer的理解写的一个代码，使用时只需实例化Positional\_Encoding类与Encoder类，使用多层Encoder时，可以设置Encoder循环的次数。

# coding = utf-8
    import torch
    import numpy as np
    import torch.nn as nn
    import torch.nn.functional as F
    
    
    # Scaled dot-product attention
    class Scaled_Dot_Product_Attention(nn.Module):
        '''Scaled Dot-Product Attention '''
    
        def __init__(self):
            super(Scaled_Dot_Product_Attention, self).__init__()
    
        def forward(self, Q, K, V, scale=None):
            '''
            Args:
                Q: [batch_size, len_Q, dim_Q]
                K: [batch_size, len_K, dim_K]
                V: [batch_size, len_V, dim_V]
                scale: 缩放因子 论文为根号dim_K
                Q, K, V 维度相同
            Return:
                self-attention后的张量，以及attention张量
            '''
            # attention = torch.matmul(Q, K.transpose(-1, -2))
            attention = torch.matmul(Q, K.permute(0, 2, 1))
            if scale:
                attention = attention * scale
            # if mask:  # TODO change this
            #     attention = attention.masked_fill_(mask == 0, -1e9)
            # attention = F.softmax(attention, dim=2)
            attention = F.softmax(attention, dim=-1)
            context = torch.matmul(attention, V)
            return context
    
    
    # multiHeadAttention
    class MultiHeadAttention(nn.Module):
        def __init__(self, model_dim, n_heads, dropout=0.0):
            """
            :param model_dim: the model dimension
            :param n_heads: head numbers
            :param dropout: dropout
            """
            super(MultiHeadAttention, self).__init__()
            self.n_heads = n_heads
            self.model_dim = model_dim
            self.dropout = nn.Dropout(dropout)
            assert self.model_dim % self.n_heads == 0
    
            self.dim_v = model_dim // n_heads
            self.W_Q = nn.Linear(model_dim, self.dim_v * n_heads, bias=False)
            self.W_K = nn.Linear(model_dim, self.dim_v * n_heads, bias=False)
            self.W_V = nn.Linear(model_dim, self.dim_v * n_heads, bias=False)
            self.attention = Scaled_Dot_Product_Attention()
            self.fc = nn.Linear(n_heads * self.dim_v, model_dim, bias=False)
            self.dropout = nn.Dropout(dropout)
            self.layer_norm = nn.LayerNorm(model_dim)
    
        def forward(self, x):
            """
            :param input_Q: [batch_size, len_Q, model_dim]
            :param input_K: [batch_size, len_K, model_dim]
            :param input_V: [batch_size, len_V=len_K, model_dim]
            :param attention_mask:[batch_size, seq_len, seq_len]
            :return:
            """
            residual, batch_size = x, x.size(0)
            Q = self.W_Q(x).view(batch_size * self.n_heads, -1, self.dim_v)
            K = self.W_K(x).view(batch_size * self.n_heads, -1, self.dim_v)
            V = self.W_V(x).view(batch_size * self.n_heads, -1, self.dim_v)
            scale = K.size(-1) ** -0.5  # 缩放因子
            context = self.attention(Q, K, V, scale)
    
            context = context.view(batch_size, -1, self.dim_v * self.n_heads)
            out = self.fc(context)
            out = self.dropout(out)
            out = out + x  # 残差连接
            out = self.layer_norm(out)
            return out
    
    
    class Position_wise_Feed_Forward(nn.Module):
        def __init__(self, model_dim, hidden, dropout=0.0):
            super(Position_wise_Feed_Forward, self).__init__()
            self.fc1 = nn.Linear(model_dim, hidden)
            self.fc2 = nn.Linear(hidden, model_dim)
            self.dropout = nn.Dropout(dropout)
            self.layer_norm = nn.LayerNorm(model_dim)
    
        def forward(self, x):
            """
            :param x:[batch_size, seq_len, model_dim]
            :return:
            """
            out = self.fc1(x)
            out = F.relu(out)
            out = self.fc2(out)
            out = self.dropout(out)
            out = out + x  # 残差连接
            out = self.layer_norm(out)
            return out
    
    
    class Encoder(nn.Module):
        def __init__(self, model_dim, n_heads, hidden, dropout):
            super(Encoder, self).__init__()
            self.attention = MultiHeadAttention(model_dim, n_heads, dropout)
            self.feed_forward = Position_wise_Feed_Forward(model_dim, hidden, dropout)
    
        def forward(self, x):
            """
            :param x: [batch_size, seq_len]
            :return:
            """
            out = self.attention(x)
            out = self.feed_forward(out)
            return out
    
    
    class Positional_Encoding(nn.Module):
        def __init__(self, embed, pad_size, dropout, device):
            super(Positional_Encoding, self).__init__()
            self.device = device
            self.pe = torch.tensor(
                [[pos / (10000.0 ** (i // 2 * 2.0 / embed)) for i in range(embed)] for pos in range(pad_size)])
            self.pe[:, 0::2] = np.sin(self.pe[:, 0::2])
            self.pe[:, 1::2] = np.cos(self.pe[:, 1::2])
            self.dropout = nn.Dropout(dropout)
    
        def forward(self, x):
            """
            :param x: [seq_len, batch_size, model_dim]
            :return:
            """
            out = x + nn.Parameter(self.pe, requires_grad=False).to(self.device)
            out = self.dropout(out)
            return out

多层Encoder的设置方法：

self.encoders = nn.ModuleList(
                [copy.deepcopy(self.encoder) for _ in range(5)]
            )

其中，self.encoder为Encoder类的实例化对象。