Transformer代码完全解读

Transformer是一种基于自注意力机制的神经网络模型，广泛应用于自然语言处理任务中，如机器翻译、文本生成等。下面是Transformer的代码解读：

首先，我们需要导入必要的库和模块：

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

接着，我们定义了一个叫做PositionalEncoding的类，用于对输入的序列进行位置编码。位置编码的目的是为了让模型能够感知输入序列中每个元素的位置信息，从而更好地处理序列中的长距离依赖关系。

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.1, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        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() * (-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).transpose(0, 1)
        self.register_buffer('pe', pe)

    def forward(self, x):
        x = x + self.pe[:x.size(0), :]
        return self.dropout(x)

在这个类中，我们首先定义了一个构造函数，其中d_model表示输入序列的维度，dropout表示dropout的概率，max_len表示输入序列的最大长度。在构造函数中，我们首先调用了父类的构造函数，然后定义了一个dropout层。

接着，我们创建了一个max_len x d_model的矩阵pe，并对其进行位置编码。具体来说，我们首先创建了一个长度为max_len的位置向量position，然后对每个位置向量应用一组不同的正弦和余弦函数，得到一个d_model维的位置编码向量。最后，我们将所有位置编码向量拼接成一个矩阵，并将其转置，以便与输入序列进行相加。

在forward函数中，我们将输入序列x与位置编码矩阵相加，并对结果进行dropout操作。

接下来，我们定义了一个叫做MultiHeadAttention的类，用于实现多头注意力机制。多头注意力机制是指将输入序列分别映射到多个不同的子空间中，并在每个子空间中计算注意力分数，最后将所有子空间的注意力分数加权求和得到最终的输出。

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, nhead, dropout=0.1):
        super(MultiHeadAttention, self).__init__()
        self.nhead = nhead
        self.d_model = d_model
        self.head_dim = d_model // nhead

        self.qkv_proj = nn.Linear(d_model, 3 * d_model)
        self.out_proj = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, query, key, value, attn_mask=None):
        batch_size = query.size(0)

        qkv = self.qkv_proj(query).chunk(3, dim=-1)
        q, k, v = qkv[0], qkv[1], qkv[2]

        q = q.view(batch_size * self.nhead, -1, self.head_dim).transpose(0, 1)
        k = k.view(batch_size * self.nhead, -1, self.head_dim).transpose(0, 1)
        v = v.view(batch_size * self.nhead, -1, self.head_dim).transpose(0, 1)

        attn_scores = torch.bmm(q, k.transpose(1, 2))
        attn_scores = attn_scores / math.sqrt(self.head_dim)

        if attn_mask is not None:
            attn_scores = attn_scores.masked_fill(attn_mask == 0, -1e9)

        attn_probs = F.softmax(attn_scores, dim=-1)
        attn_probs = self.dropout(attn_probs)

        attn_output = torch.bmm(attn_probs, v)
        attn_output = attn_output.transpose(0, 1).contiguous().view(batch_size, -1, self.d_model)

        attn_output = self.out_proj(attn_output)
        attn_output = self.dropout(attn_output)

        return attn_output

在这个类中，我们首先定义了一个构造函数，其中d_model表示输入序列的维度，nhead表示头的数量，dropout表示dropout的概率。在构造函数中，我们首先调用了父类的构造函数，然后定义了一个线性层qkv_proj，用于将输入序列映射到三个不同的子空间中。接着，我们定义了一个线性层out_proj，用于将多头注意力机制的输出映射回原始的输入维度。最后，我们定义了一个dropout层。

在forward函数中，我们首先获取输入序列的batch_size，并将输入序列通过线性层qkv_proj映射到三个不同的子空间中。然后，我们将每个子空间的向量分别重塑为(batch_size * nhead, seq_len, head_dim)的形状，并将其转置，以便进行矩阵乘法。接着，我们计算每个位置之间的注意力分数，并对其进行缩放。如果存在attn_mask，则将其应用于注意力分数。然后，我们对注意力分数进行softmax操作，并对结果进行dropout。接着，我们将注意力分数与value矩阵相乘，并将结果重塑为(batch_size, seq_len, d_model)的形状。最后，我们将输出通过线性层out_proj映射回原始的输入维度，并对结果进行dropout。

最后，我们定义了一个叫做TransformerEncoderLayer的类，用于实现Transformer的编码器层。编码器层由两个子层组成：多头自注意力机制和前馈神经网络。

class TransformerEncoderLayer(nn.Module):
    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1):
        super(TransformerEncoderLayer, self).__init__()
        self.self_attn = MultiHeadAttention(d_model, nhead, dropout=dropout)
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(p=dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)

    def forward(self, src, src_mask=None):
        src2 = self.self_attn(src, src, src, attn_mask=src_mask)
        src = src + self.dropout(src2)
        src = self.norm1(src)

        src2 = self.linear2(self.dropout(F.relu(self.linear1(src))))
        src = src + self.dropout(src2)
        src = self.norm2(src)

        return src

在这个类中，我们首先定义了一个构造函数，其中d_model表示输入序列的维度，nhead表示头的数量，dim_feedforward表示前馈神经网络的隐藏层维度，dropout表示dropout的概率。在构造函数中，我们定义了一个多头自注意力机制self_attn，一个线性层linear1，一个dropout层，一个线性层linear2，以及两个LayerNorm层。

在forward函数中，我们首先使用多头自注意力机制self_attn对输入序列进行编码，并将结果与原始输入序列相加。然后，我们对结果进行归一化，并通过一个前馈神经网络进行非线性变换。最后，我们再次将结果与原始输入序列相加，并对结果进行归一化。