transformer代码

以下是一个简单的Transformer模型的代码实现，包括Encoder和Decoder部分：

import torch.nn as nn
import torch.nn.functional as F
import torch

class MultiHeadAttention(nn.Module):
    def __init__(self, n_heads, d_model, dropout):
        super(MultiHeadAttention, self).__init__()
        self.n_heads = n_heads
        self.d_model = d_model
        
        self.key_layer = nn.Linear(d_model, d_model*n_heads)
        self.query_layer = nn.Linear(d_model, d_model*n_heads)
        self.value_layer = nn.Linear(d_model, d_model*n_heads)
        self.dropout = nn.Dropout(dropout)
        self.out = nn.Linear(d_model*n_heads, d_model)
        
    def forward(self, queries, keys, values, mask=None):
        bsz = queries.size(0)

        queries = self.query_layer(queries).view(bsz, -1, self.n_heads, self.d_model).transpose(1,2)
        keys = self.key_layer(keys).view(bsz, -1, self.n_heads, self.d_model).transpose(1,2)
        values = self.value_layer(values).view(bsz, -1, self.n_heads, self.d_model).transpose(1,2)
        
        attn_scores = torch.matmul(queries, keys.transpose(-2,-1))
        attn_scores = attn_scores / torch.sqrt(torch.tensor(self.d_model, dtype=torch.float32))
        if mask is not None:
            attn_scores = attn_scores.masked_fill(mask == 0, -1e9)
        attn_scores = F.softmax(attn_scores, dim=-1)
        attn_scores = self.dropout(attn_scores)
        attn_output = torch.matmul(attn_scores, values)
        attn_output = attn_output.transpose(1,2).contiguous().view(bsz, -1, self.n_heads*self.d_model)
            
        return self.out(attn_output)
        

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_seq_len):
        super(PositionalEncoding, self).__init__()
        self.d_model = d_model
        pe = torch.zeros(max_seq_len, d_model)
        position = torch.arange(0, max_seq_len, dtype=torch.float).unsqueeze(1)
        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):
        return x + self.pe[:, :x.size(1)].detach()

        
class FeedForward(nn.Module):
    def __init__(self, d_model, d_ff, dropout):
        super(FeedForward, self).__init__()
        self.linear1 = nn.Linear(d_model, d_ff)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(d_ff, d_model)
        
    def forward(self, x):
        x = F.relu(self.linear1(x))
        x = self.dropout(x)
        x = self.linear2(x)
        return x
    
        
class EncoderLayer(nn.Module):
    def __init__(self, n_heads, d_model, d_ff, dropout):
        super(EncoderLayer, self).__init__()
        self.self_attn = MultiHeadAttention(n_heads, d_model, dropout)
        self.ff = FeedForward(d_model, d_ff, dropout)
        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, x, mask=None):
        attn_output = self.self_attn(x, x, x, mask)
        x = x + self.dropout1(attn_output)
        x = self.norm1(x)
        ff_output = self.ff(x)
        x = x + self.dropout2(ff_output)
        x = self.norm2(x)
        return x

    
class DecoderLayer(nn.Module):
    def __init__(self, n_heads, d_model, d_ff, dropout):
        super(DecoderLayer, self).__init__()
        self.self_attn = MultiHeadAttention(n_heads, d_model, dropout)
        self.src_attn = MultiHeadAttention(n_heads, d_model, dropout)
        self.ff = FeedForward(d_model, d_ff, dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)
        
    def forward(self, x, memory, src_mask=None, tgt_mask=None):
        self_attn_output = self.self_attn(x, x, x, tgt_mask)
        x = x + self.dropout1(self_attn_output)
        x = self.norm1(x)
        src_attn_output = self.src_attn(x, memory, memory, src_mask)
        x = x + self.dropout2(src_attn_output)
        x = self.norm2(x)
        ff_output = self.ff(x)
        x = x + self.dropout3(ff_output)
        x = self.norm3(x)
        return x


class Transformer(nn.Module):
    def __init__(self, src_vocab_size, tgt_vocab_size, max_seq_len, d_model, n_heads, n_layers, d_ff, dropout):
        super(Transformer, self).__init__()
        self.src_embedding = nn.Embedding(src_vocab_size, d_model)
        self.tgt_embedding = nn.Embedding(tgt_vocab_size, d_model)
        self.pos_encoder = PositionalEncoding(d_model, max_seq_len)
        self.encoder_layers = nn.ModuleList([EncoderLayer(n_heads, d_model, d_ff, dropout) for _ in range(n_layers)])
        self.decoder_layers = nn.ModuleList([DecoderLayer(n_heads, d_model, d_ff, dropout) for _ in range(n_layers)])
        self.out = nn.Linear(d_model, tgt_vocab_size)
        
    def forward(self, src, tgt, src_mask=None, tgt_mask=None):
        src_emb = self.src_embedding(src)
        tgt_emb = self.tgt_embedding(tgt)
        src_emb = self.pos_encoder(src_emb)
        tgt_emb = self.pos_encoder(tgt_emb)
        memory = self.encode(src_emb, src_mask)
        output = self.decode(tgt_emb, memory, src_mask, tgt_mask)
        return F.log_softmax(self.out(output), dim=-1)
    
    def encode(self, x, mask=None):
        for layer in self.encoder_layers:
            x = layer(x, mask)
        return x
    
    def decode(self, x, memory, src_mask=None, tgt_mask=None):
        for layer in self.decoder_layers:
            x = layer(x, memory, src_mask, tgt_mask)
        return x

该模型包含了Multi-Head Attention、Positional Encoding、FeedForward、Encoder Layer、Decoder Layer以及Transformer本身的实现。其中，Multi-Head Attention和FeedForward是Transformer中的两个基本模块，Encoder和Decoder部分分别由多个Encoder Layer和Decoder Layer组成。Transformer模型的正向传播过程是先对源语言进行编码，得到一个中间语言表示，然后再对目标语言进行解码，最终输出目标语言的概率分布。