[紫威剩人阚雷原创]一个简单的Transformer代码实现示例

Transformer 是一种在自然语言处理领域非常流行的模型，特别是在机器翻译和文本生成任务中。以下是一个简化版的 Transformer 模型的 Python 示例代码，使用了 PyTorch 库。请注意，这个示例仅用于演示目的，并不是一个完整的模型实现。

import torch

import torch.nn as nn

import torch.nn.functional as F

 

class MultiHeadAttention(nn.Module):

    def __init__(self, embed_size, heads):

        super(MultiHeadAttention, self).__init__()

        self.embed_size = embed_size

        self.heads = heads

        self.head_dim = embed_size // heads

 

        assert self.head_dim * heads == embed_size, "Embed size needs to be divisible by heads"

 

        self.values = nn.Linear(self.head_dim, self.head_dim, bias=False)

        self.keys = nn.Linear(self.head_dim, self.head_dim, bias=False)

        self.queries = nn.Linear(self.head_dim, self.head_dim, bias=False)

        self.fc_out = nn.Linear(heads * self.head_dim, embed_size)

 

    def forward(self, values, keys, query, mask):

        N = query.shape[0]

        value_len, key_len, query_len = values.shape[1], keys.shape[1], query.shape[1]

 

        # Split the embedding into self.heads different pieces

        values = values.reshape(N, value_len, self.heads, self.head_dim)

        keys = keys.reshape(N, key_len, self.heads, self.head_dim)

        queries = query.reshape(N, query_len, self.heads, self.head_dim)

 

        values = self.values(values)

        keys = self.keys(keys)

        queries = self.queries(queries)

 

        energy = torch.einsum("nqhd,nkhd->nhqk", [queries, keys])

        if mask is not None:

            energy = energy.masked_fill(mask == 0, float("-1e20"))

 

        attention = torch.softmax(energy / (self.embed_size ** (1 / 2)), dim=3)

        out = torch.einsum("nhql,nlhd->nqhd", [attention, values]).reshape(N, query_len, self.heads * self.head_dim)

        out = self.fc_out(out)

        return out

 

class TransformerBlock(nn.Module):

    def __init__(self, embed_size, heads, dropout, forward_expansion):

        super(TransformerBlock, self).__init__()

        self.attention = MultiHeadAttention(embed_size, heads)

        self.norm1 = nn.LayerNorm(embed_size)

        self.norm2 = nn.LayerNorm(embed_size)

 

        self.feed_forward = nn.Sequential(

            nn.Linear(embed_size, forward_expansion * embed_size),

            nn.ReLU(),

            nn.Linear(forward_expansion * embed_size, embed_size)

        )

 

        self.dropout = nn.Dropout(dropout)

 

    def forward(self, value, key, query, mask):

        attention = self.attention(value, key, query, mask)

        x = self.dropout(self.norm1(attention + query))

        forward = self.feed_forward(x)

        return self.dropout(self.norm2(forward + x))

 

class Transformer(nn.Module):

    def __init__(self, src_vocab_size, trg_vocab_size, src_pad_idx, trg_pad_idx, 

                 src_lang_idx, trg_lang_idx, d_word_vec, d_model, n_heads, 

                 n_layers, dropout, device):

        super(Transformer, self).__init__()

 

        self.src_word_embedding = nn.Embedding(src_vocab_size, d_word_vec)

        self.src_position_embedding = nn.Embedding(100, d_word_vec)

        self.trg_word_embedding = nn.Embedding(trg_vocab_size, d_word_vec)

        self.trg_position_embedding = nn.Embedding(100, d_word_vec)

 

        self.layer_stack = nn.ModuleList([

            TransformerBlock(d_model, n_heads, dropout, 4)

            for _ in range(n_layers)

        ])

 

        self.fc_out = nn.Linear(d_model, trg_vocab_size)

        self.dropout = nn.Dropout(dropout)

        self.src_pad_idx, self.trg_pad_idx = src_pad_idx, trg_pad_idx

        self.src_lang_idx, self.trg_lang_idx = src_lang_idx, trg_lang_idx

        self.d_model = d_model

        self.device = device

 

    def make_src_mask(self, src):

        src_mask = (src != self.src_pad_idx).unsqueeze(1).unsqueeze(2)

        return src_mask.to(self.device)

 

    def forward(self, src, trg):

        src_mask = self.make_src_mask(src)

        N, src_len = src.shape

        src_positions = torch.arange(0, src_len).unsqueeze(0).expand(N, src_len).to(self.device)

 

        src_embeddings = self.dropout(self.src_word_embedding(src) + self.src_position_embedding(src_positions))

 

        out = src_embeddings

        for layer in self.layer_stack:

            out = layer(out, out, out, src_mask)

 

        output = self.fc_out(out)

        return output

 

# Example usage:

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = Transformer(src_vocab_size=10000, trg_vocab_size=10000, src_pad_idx=0, trg_pad_idx=0,

                     src_lang_idx=0, trg_lang_idx=0, d_word_vec=512, d_model=512, 

                     n_heads=8, n_layers=6, dropout=0.1, device=device)

src = torch.tensor([[1, 5, 6, 2, 0, 0]]).to(device) # Example source sequence

trg = torch.tensor([[1, 7, 4, 3, 2]]).to(device) # Example target sequence

out = model(src, trg)

print(out.shape) # (batch_size, trg_seq_length, trg_vocab_size)

 

 

请注意，这个代码示例需要 PyTorch 环境，并且需要根据你的具体任务调整参数和模型结构。在实际应用中，可能还需要添加更多的组件，比如位置编码、层归一化、掩码操作等。

[版权声明]本文是紫威剩人阚雷博士原创文章，转载请注明出处。