【机器学习】一文搞懂算法模型之：Transformer

    # -*- coding:utf-8 -*-
    # @Time   : 2024-03-17
    # @Author : Carl_DJ
    
    '''
    import torch  
    import torch.nn as nn  
    import torch.nn.functional as F  
      
    # 定义位置编码类  
    class PositionalEncoding(nn.Module):  
    def __init__(self, d_model, max_len=5000):  
        super(PositionalEncoding, self).__init__()  
          
        # 计算位置编码  
        pe = torch.zeros(max_len, d_model)  
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)  
        div_term = torch.exp(torch.arange(0, d_model, 2).float() *  
                                 -(np.log(torch.tensor(10000.0)) / d_model))  
        pe[:, 0::2] = torch.sin(position * div_term)  
        pe[:, 1::2] = torch.cos(position * div_term)  
        pe = pe.unsqueeze(0).transpose(0, 1)  
        self.register_buffer('pe', pe)  
          
    def forward(self, x):  
        x = x + self.pe[:, :x.size(1)]  
        return x  
      
    # 定义Transformer编码器层  
    class TransformerEncoderLayer(nn.Module):  
    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1):  
        super(TransformerEncoderLayer, self).__init__()  
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)  
        self.linear1 = nn.Linear(d_model, dim_feedforward)  
        self.dropout = nn.Dropout(dropout)  
        self.linear2 = nn.Linear(dim_feedforward, d_model)  
          
        self.norm1 = nn.LayerNorm(d_model)  
        self.norm2 = nn.LayerNorm(d_model)  
        self.dropout1 = nn.Dropout(dropout)  
        self.dropout2 = nn.Dropout(dropout)  
          
    def forward(self, src, src_mask=None, src_key_padding_mask=None):  
        # 自注意力机制  
        src2, attn = self.self_attn(src, src, src, attn_mask=src_mask,  
                                      key_padding_mask=src_key_padding_mask)  
        src = src + self.dropout1(src2)  
        src = self.norm1(src)  
          
        # 前馈神经网络  
        src2 = self.linear2(self.dropout(F.relu(self.linear1(src))))  
        src = src + self.dropout2(src2)  
        src = self.norm2(src)  
          
        return src, attn  
      
    # 定义Transformer编码器  
    class TransformerEncoder(nn.Module):  
    def __init__(self, encoder_layer, num_layers, norm=None):  
        super(TransformerEncoder, self).__init__()  
        self.layers = nn.ModuleList([encoder_layer for _ in range(num_layers)])  
        self.num_layers = num_layers  
        self.norm = norm  
          
    def forward(self, src, mask=None, src_key_padding_mask=None):  
        output = src  
          
        # 堆叠多个编码器层  
        for mod in self.layers:  
            output, attn = mod(output, mask, src_key_padding_mask)  
          
        if self.norm:  
            output = self.norm(output)  
          
        return output  
      
    # 定义完整的Transformer模型  
    class TransformerModel(nn.Module):  
    def __init__(self, src_vocab_size, d_model, nhead, num_encoder_layers, dim_feedforward, max_position_embeddings):  
        super(TransformerModel, self).__init__()  
        self.src_mask = None  
        self.position_enc = PositionalEncoding(d_model, max_position_embeddings)  
        encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout=0.1)  
        encoder_norm = nn.LayerNorm(d_model)  
        self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, norm=encoder_norm)  
    		self.src_embedding = nn.Embedding(src_vocab_size, d_model)  
    	self.d_model = d_model  
      
    def generate_square_subsequent_mask(self, sz):  
    mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)  
    mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))  
    return mask  
      
    def forward(self, src, src_mask=None):  
    if src_mask is None:  
        device = src.device  
        bsz, seq_len = src.size()  
        src_mask = self.generate_square_subsequent_mask(seq_len).to(device)  
      
    src = self.src_embedding(src) * math.sqrt(self.d_model)  
    src = self.position_enc(src)  
    output = self.encoder(src, src_mask)  
    output = self.decoder(output)  
      
    return output
    
    
    
    
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
    
    代码解读