大模型原理:构建GPT-2大语言模型-实现层归一化类

准备

import torch
import torch.nn as nn

torch.manual_seed(123)
# 生成2行5列张量,作为两个训练样本,每个样本包含5个维度(特征)
batch_example = torch.randn(2, 5)
# 神经网络层包括一个线性层和一个非线性激活函数ReLU
layer = nn.Sequential(nn.Linear(5, 6), nn.ReLU())
# 执行计算
out = layer(batch_example)
print(out)

层归一化类

class LayerNorm(nn.Module):
    def __init__(self, emb_dim):
        """

        :param emb_dim:嵌入维度
        """
        super().__init__()
        # eps是一个小常数(epsilon),在归一化过程中会被加到方差上以防止除零错误
        self.eps = 1e-5
        """
        scale和shift是两个可训练的参数(与输入维度相同),
        在训练过程中发现调整它们可以改善模型的训练任务表现,那么大语言模型会自动进行调整
        """
        self.scale = nn.Parameter(torch.ones(emb_dim))
        self.shift = nn.Parameter(torch.zeros(emb_dim))

    def forward(self, x):
        mean = x.mean(dim=-1, keepdim=True)
        """
        设置unbiased=False,在方差计算中,会使用样本数量 n 作为方差公式的除数。
        这种方法没有使用 n-1 作为分母,以调整偏差。
        因此,这种方法得到的是所谓有偏方差估计。
        对于嵌入维度 n 非常大的大语言模型(如GPT-2),使用 n 和 n-1 的差异在实际中几乎可以忽略。
        选择这种方法是为了确保与GPT-2模型的归一化层兼容,因为原始GPT-2模型是用TensorFlow实现的,TensorFlow默认行为是 n
        """
        var = x.var(dim=-1, keepdim=True, unbiased=False)
        norm_x = (x - mean) / torch.sqrt(var + self.eps)
        # 相当于 wx + b
        return self.scale * norm_x + self.shift

执行

# 实例化层归一化对象
ln = LayerNorm(emb_dim=5)
# 计算归一化
out_ln = ln(batch_example)
# 沿着最后一个维度(列)计算均值,且输出维度与输入维度相同
mean = out_ln.mean(dim=-1, keepdim=True)
# 沿着最后一个维度(列)计算方差,且输出维度与输入维度相同
var = out_ln.var(dim=-1, unbiased=False, keepdim=True)
print("层归一化结果:\n", out_ln)

# 关闭科学记数法
torch.set_printoptions(sci_mode=False)
print("均值:\n", mean)
print("方差:\n", var)

执行结果:

层归一化结果:
 tensor([[ 0.5528,  1.0693, -0.0223,  0.2656, -1.8654],
        [ 0.9087, -1.3767, -0.9564,  1.1304,  0.2940]], grad_fn=<AddBackward0>)
均值:
 tensor([[-0.0000],
        [ 0.0000]], grad_fn=<MeanBackward1>)
方差:
 tensor([[1.0000],
        [1.0000]], grad_fn=<VarBackward0>)

结果表明,层归一化代码正常工作,成功地将两个输入的值归一化,使其均值为0,方差为1。

总结

现在实现GPT架构所需的两个关键构建块:GPT主干和层归一化,如下图:

接下来,将研究 GELU激活函数,它是大语言模型中使用的激活函数之一,而不是前面用过的传统ReLU函数。


原文出处:/show_1GW2eJda5VjE.html