1、前馈层

前馈层接收自注意力层的输出作为输入。

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

class FeedForward(nn.Module):
    def __init__(self, d_model=512, d_ff=2048, dropout=0.1):
        super().__init__()

        # d_ff 默认设置为2048
        self.linear_1 = nn.Linear(d_model, d_ff)
        self.dropout = nn.Dropout(dropout)
        self.linear_2 = nn.Linear(d_ff, d_model)

    def forward(self, x):
        x = self.linear_1(x)
        x = self.dropout(F.relu(x))
        x = self.linear_2(x)

        return x

2、残差连接

残差连接和层归一化技术，进一步提升训练的稳定性。
残差连接主要是指使用一条直连通道直接将对应子层的输入连接到输出，避免在优化过程中因网络过深而产生潜在的梯度消失问题。
x = f(x) + x

3、层归一化

为了使每一层的输入/输出稳定在一个合理的范围内，层归一化技术被进一步引入每个Transformer快中：
LN(x) = a (x-u)/c + b
其中u和c分别表示均值和方差，用于将数据平移缩放到均值为0、方差为1的标准分布，a和b是可学习的参数。层归一化技术可以有效地缓解优化过程中潜在的不稳定、收敛速度慢等问题。

import torch
from torch import nn

class Norm(nn.Module):
    def __init__(self, d_model, eps=1e-6):
        super(Norm, self).__init__()

        self.size = d_model

        # 层归一化包含两个可以学习的参数
        self.alpha = nn.Parameter(torch.ones(self.size))
        self.bias = nn.Parameter(torch.zeros(self.size))

        self.eps = eps

    def forward(self, x):
        norm = self.alpha * (x - x.mean(dim=-1, keepdim=True)) / (x.std(dim=-1, keepdim=True) + self.eps) + self.bias
        return norm