引言

本文開始從零實(shí)現(xiàn)GPT1做一個(gè)小說續(xù)寫器，即只需要給出一些文本，讓模型幫你續(xù)寫，主要內(nèi)容包含：

• 模型編寫
• 訓(xùn)練適配小說的中文分詞器
• 將小說按固定大小拆分生成數(shù)據(jù)集
• 拆分訓(xùn)練/測(cè)試集
• 訓(xùn)練
• 體驗(yàn)小說續(xù)寫效果

同時(shí)結(jié)合HuggingFace的transformers，可以將處理好的數(shù)據(jù)集、訓(xùn)練好的分詞器和模型上傳到HuggingFace Hub。

本文主要實(shí)現(xiàn)模型編寫，剩下的內(nèi)容請(qǐng)見下篇文章。

模型架構(gòu)

GPT模型架構(gòu)如上圖所示，由多層Tranformer Decoder組成的單向語言模型，是Tranformer的一個(gè)變種。它的Transformer Block比較簡(jiǎn)單，由兩個(gè)子層組成，第一個(gè)子層輸入上應(yīng)用一個(gè)多頭注意力層，輸入和輸出經(jīng)過殘差連接，緊著的是一個(gè)層歸一化；第二個(gè)子層是前饋層、殘差連接和層歸一化。

整個(gè)GPT可以分為三部分：

1. 輸入層
2. 編碼層
3. 輸出層

輸入層計(jì)算出Transformer Block的輸入表示；編碼層經(jīng)過疊加的多層Transformer Block進(jìn)行編碼；最后輸出層應(yīng)用Softmax計(jì)算輸出標(biāo)記的分布。

其訓(xùn)練過程包含兩個(gè)階段：無監(jiān)督預(yù)訓(xùn)練和有監(jiān)督微調(diào)。

無監(jiān)督階段可以在大規(guī)模文本語料上學(xué)習(xí)一個(gè)高容量的語言模型，然后可以根據(jù)下游具體任務(wù)進(jìn)行微調(diào)。

無監(jiān)督預(yù)訓(xùn)練

GPT是一個(gè)單向模型，也是僅解碼器模型(Decoder Only)，即只能自左向右(或反之)對(duì)文本序列建模，采用的是Transformer的解碼器結(jié)構(gòu)，同時(shí)引入了同樣的解碼策略保證輸入文本每個(gè)位置只能依賴當(dāng)前和過去時(shí)刻的信息。

給定文本序列$w=w_1{w}_2?w_n$，首先通過輸入層將其編碼成稠密向量：
$${u\text{u}u}_i={u\text{u}u}_i^e+{u\text{u}u}_i^p\text{\hspace{1em}(1)}$$
輸入層由兩個(gè)子層組成：詞嵌入層和位置編碼層。

其中${u\text{u}u}_i^e$是$w_i$經(jīng)過詞嵌入層得到的詞向量；${u\text{u}u}_i^p$是$w_i$的經(jīng)過位置編碼層得到的位置向量；${u\text{u}u}_i$為第$i$個(gè)位置的標(biāo)記經(jīng)過輸入層后的輸出。

GPT的位置編碼和原始Transformer中固定的不同，它是一種可學(xué)習(xí)的位置編碼。

經(jīng)過輸入層得到每個(gè)標(biāo)記帶位置信息的詞嵌入表示序列$u\text{u}u={u\text{u}u}_1?{u\text{u}u}_n$，接著將$u\text{u}u$輸入GPT的編碼層，編碼層由$L$個(gè)Transformer Block組成，每一層的Block都能計(jì)算出帶有上下文信息的向量表示，經(jīng)過多層編碼后，能得到更復(fù)雜、強(qiáng)大的向量表示，計(jì)算過程為：
$${h\text{h}h}^l={\text{transformer\_block}}^l({h\text{h}h}^{l?1})?l\in [1,L]\text{\hspace{1em}(2)}$$
其中我們令${h\text{h}h}^0=u\text{u}u$，即輸入層計(jì)算出來的輸出；${h\text{h}h}^l\in {\mathbb{R}}^{d\times n}$表示由第$l$層計(jì)算出來的表示向量序列，$d$是模型隱藏層維度，$n$為序列長(zhǎng)度；$L$為總層數(shù)。

而輸出層基于最后一層的向量表示${h\text{h}h}^L$計(jì)算每個(gè)位置上輸出標(biāo)記的概率分布：
$$P(w_i|w_1,?,w_{i?1})=\text{softmax}({W\text{W}W}^e{h\text{h}h}_i^L)\text{\hspace{1em}(3)}$$
這里${W\text{W}W}^e\in {\mathbb{R}}^{|\mathbb{V}|\times d}$是詞向量矩陣；$|\mathbb{V}|$為詞表大??；注意這里${h\text{h}h}_i^L$的維度是$d\times 1$。

然后使用一個(gè)常規(guī)的語言建模目標(biāo)優(yōu)化$w$的最大似然估計(jì)：
$${\mathcal{L}}^{\text{PT}}=?\sum _i\log P(w_i|w_{i?k}?w_{i?1};\Theta )\text{\hspace{1em}(4)}$$
這里的$k$是上下文窗口，根據(jù)前$k$個(gè)標(biāo)記來預(yù)測(cè)當(dāng)前標(biāo)記；$\Theta$表示模型參數(shù)。

這就是預(yù)訓(xùn)練(pretrain)階段的損失函數(shù)。

有監(jiān)督微調(diào)

無監(jiān)督預(yù)訓(xùn)練使得模型具有一定的通用語義表示能力，下游任務(wù)微調(diào)目的使通用語義表示可以適配不同具體的下游任務(wù)。

微調(diào)一般需要利用有標(biāo)簽數(shù)據(jù)集進(jìn)行，假設(shè)一個(gè)有標(biāo)簽數(shù)據(jù)集$\mathcal{C}$，其中每個(gè)樣本包含一個(gè)輸入序列$x=x_1{x}_2?x_n$和一個(gè)輸出標(biāo)簽$y$。

將$x$輸入給預(yù)訓(xùn)練好的模型，我們用最后一層Transformer Block的最后一個(gè)位置的輸出${h\text{h}h}_n^L$來進(jìn)行預(yù)測(cè)，具體地可以接一個(gè)全連接層結(jié)合$\text{softmax}$函數(shù)得到預(yù)測(cè)標(biāo)簽的概率分布：
$$p(y|x_1?x_n)=\text{softmax}({h\text{h}h}_n^L{W\text{W}W}^y)\text{\hspace{1em}(5)}$$
其中${W\text{W}W}^y\in {\mathbb{R}}^{d\times c}$為全連接層參數(shù)；$c$為標(biāo)簽個(gè)數(shù)。通過對(duì)整個(gè)標(biāo)注數(shù)據(jù)集進(jìn)行優(yōu)化，我們又可以得到微調(diào)目標(biāo)函數(shù)：
$${\mathcal{L}}^{\text{FT}}(\mathcal{C})=?\sum _{(x,y)}\log P(y|x_1?x_n)\text{\hspace{1em}(6)}$$
在下游任務(wù)微調(diào)過程中，如果僅針對(duì)微調(diào)目標(biāo)進(jìn)行優(yōu)化，很可能會(huì)使模型遺忘預(yù)訓(xùn)練階段所學(xué)習(xí)到的通用語義表示知識(shí)，從而損失模型的通用性和泛化能力，即災(zāi)難性遺忘(Catastrophic Forgetting)。因此將語言建模任務(wù)作為一個(gè)輔助目標(biāo)函數(shù)加到微調(diào)階段可以有助于學(xué)習(xí)，具體地，我們優(yōu)化下面的目標(biāo)函數(shù)：
$$\mathcal{L}={\mathcal{L}}^{\text{FT}}(\mathcal{C})+\lambda {\mathcal{L}}^{\text{PT}}(\mathcal{C})\text{\hspace{1em}(7)}$$
其中$\lambda$是用于平衡這兩個(gè)目標(biāo)函數(shù)的權(quán)重，可以取值$0.5$。

模型實(shí)現(xiàn)

本節(jié)我們開始從零實(shí)現(xiàn)GPT，有了上篇文章從零實(shí)現(xiàn)Transformer的基礎(chǔ)，實(shí)現(xiàn)GPT也不是太難。

本次實(shí)現(xiàn)參考了HuggingFace的源碼，使得我們后面可以很容易的應(yīng)用HuggingFace實(shí)現(xiàn)的GPT。

開始之前，我們回顧下GPT論文中實(shí)現(xiàn)細(xì)節(jié)。

實(shí)現(xiàn)細(xì)節(jié)

模型設(shè)定

• 模型主要沿用原始的Transformer；
• 訓(xùn)練了一個(gè)帶掩碼自注意力頭(狀態(tài)維度768,12個(gè)頭)的12層僅解碼器的Transformer；
• 對(duì)于位置感知的前饋網(wǎng)絡(luò)，使用3072作為內(nèi)部隱狀態(tài)維度；
• 使用Adam優(yōu)化器和最大學(xué)習(xí)率2.5e-4；
• 學(xué)習(xí)率在前2000步內(nèi)逐漸從0開始線性地增加，然后使用余弦調(diào)度器降低到0；
• 在批大小為64的長(zhǎng)度為512的序列樣本上訓(xùn)練；
• 由于模型中廣泛使用層歸一化，因此簡(jiǎn)單地(高斯)權(quán)重初始化；
• 使用了一個(gè)包含40000個(gè)合并的字節(jié)對(duì)編碼(BPE)詞表；
• 應(yīng)用殘差、嵌入和注意力的Dropout為0.1進(jìn)行正則化；
• 采用了修改版的L2正則化；
• 對(duì)所有非偏置或增益權(quán)重使用$w=0.01$；
• 對(duì)于激活函數(shù)，使用GELU；
• 使用了學(xué)習(xí)的位置嵌入，而不是原始工作中的正弦版本。

微調(diào)細(xì)節(jié)

• 基本重復(fù)使用了無監(jiān)督預(yù)訓(xùn)練的超參數(shù)設(shè)置；
• 在分類器中添加了0.1的Dropout；
• 對(duì)于大多數(shù)任務(wù)，使用6.25e-5的學(xué)習(xí)率和32的批量大小；
• 模型可以快速微調(diào)，大多數(shù)情況下3個(gè)epoch就足夠了；
• 使用線性學(xué)習(xí)率衰減調(diào)度，并在0.2%的訓(xùn)練期上進(jìn)行預(yù)熱；
• 兩個(gè)損失函數(shù)間的$\lambda$設(shè)置為0.5；

我們按照從下至上的原則依次實(shí)現(xiàn)。

輸入層

上面我們知道，輸入層由兩個(gè)子層：詞嵌入層和可學(xué)習(xí)的位置編碼層組成，那就非常簡(jiǎn)單了，實(shí)際上就是兩個(gè)嵌入層：

te=nn.Embedding(vocab_size, embed_dim )  # token emebedding 詞嵌入層
pe=nn.Embedding(max_positions, embed_dim ) # 位置編碼層

vocab_size是詞表大?。?code>embed_dim是模型嵌入大??；max_positions是最大可學(xué)習(xí)位置長(zhǎng)度。

編碼層

編碼層由$L$層Transformer Block組成，每個(gè)Block的結(jié)構(gòu)如上圖所示。我們依次實(shí)現(xiàn)。

GELU

激活函數(shù)使用GELU而不是RELU，我們來看下GELU的圖像(藍(lán)線)：

其近似公式為：
$$0.5x(1+\tanh [\sqrt{2/\pi }(x+0.044715x^3)])\text{\hspace{1em}(8)}$$
從圖像可以看到，相比RELU和ELU，GELU有以下優(yōu)勢(shì)：

• 平滑性： GELU函數(shù)在整個(gè)輸入范圍內(nèi)是光滑的，而ReLU在負(fù)數(shù)部分不是光滑的(其導(dǎo)數(shù)為0)，雖然ELU在負(fù)數(shù)部分是光滑的，但變化不夠平滑。這使得GELU更容易優(yōu)化；
• 高性能： GELU函數(shù)表現(xiàn)出比ReLU和ELU更好的性能；
• 非線性：GELU函數(shù)是非線性的，引入類似sigmoid函數(shù)的變換，使得GELU函數(shù)的輸出可以落在一個(gè)更廣的范圍內(nèi)，有助于加速模型的收斂；

按照公式實(shí)現(xiàn)即可：

class GELU(nn.Module):
    def forward(self, x: Tensor) -> Tensor:
        return (
            0.5
            * x
            * (
                1.0
                + torch.tanh(
                    math.sqrt(2.0 / math.pi)
                    * (input + 0.044715 * torch.pow(input, 3.0))
                )
            )
        )

但是為了速度快一點(diǎn)，我們應(yīng)用Pytorch內(nèi)建的torch.nn.functional.gelu。

一維卷積層

OpenAI GPT的作者把Transformer中的線性層命名為一維卷積，因?yàn)樗鼈兊牟僮魇窍嗟鹊?卷積的filter大小為1)。

我們通過圖片來直觀理解一下， https://ezyang.github.io/convolution-visualizer/ 提供了一個(gè)很好地可視化頁面。

實(shí)際上filter大小為1的一維卷積就是讓輸入中每個(gè)位置與權(quán)重相乘(即序列長(zhǎng)度維度上是并行獨(dú)立計(jì)算的)，通過out_channels控制輸出維度。

我們可以通過代碼驗(yàn)證一下：

import torch
import torch.nn as nn

embed_dim = 10
seq_len = 3
batch_size = 2
hidden_size = 5
# 定義輸入數(shù)據(jù)，表示
x = torch.randn(batch_size, seq_len, embed_dim)

# 定義前饋網(wǎng)絡(luò)
fc = torch.nn.Linear(embed_dim, hidden_size)

# 定義一維卷積核
conv = torch.nn.Conv1d(embed_dim, hidden_size, kernel_size=1)

# 設(shè)置前饋網(wǎng)絡(luò)和一維卷積核的參數(shù)相同
conv.weight = nn.Parameter(fc.weight.reshape(hidden_size, embed_dim, 1))
conv.bias = fc.bias

# 計(jì)算前饋網(wǎng)絡(luò)和一維卷積的輸出結(jié)果
fc_output = fc(x)
x_conv = x.permute(0, 2, 1)
conv_output = conv(x_conv)

# 比較輸出結(jié)果是否相同
conv_output = conv_output.permute(0, 2, 1)

print(torch.allclose(fc_output, conv_output))

True

所以它只是一個(gè)命名上的技巧，實(shí)際上實(shí)現(xiàn)起來還是通過前饋網(wǎng)絡(luò)，不過與FeedForward中權(quán)重參數(shù)的維度位置相反，我們先看這里Conv1D的實(shí)現(xiàn)：

class Conv1D(nn.Module):
    def __init__(self, in_features: int, out_features: int) -> None:
        """1D-convolutional layer as defined by Radford et al. for OpenAI GPT.

        Args:
            in_features (int): the number of input features.
            out_features (int): the number of output features.
        """
        super().__init__()
        self.out_features = out_features
        self.weight = nn.Parameter(torch.empty(in_features, out_features))
        self.bias = nn.Parameter(torch.zeros(out_features))
        nn.init.normal_(self.weight, std=0.02)

    def forward(self, x: Tensor) -> Tensor:
        """

        Args:
            x (Tensor): (batch_size, seq_len, embed_dim)

        Returns:
            Tensor: (batch_size, seq_len, out_features)
        """
        # size_out (batch_size, seq_len, out_features)
        size_out = x.size()[:-1] + (self.out_features,)
        # self.bias + x @ self.weight
        # x -view-> (batch_size *  seq_len,embed_dim)
        # (batch_size * seq_len,embed_dim) x (embed_dim, out_features)
        # -> (batch_size * seq_len, out_features)
        x = torch.addmm(self.bias, x.view(-1, x.size(-1)), self.weight)
        # x (batch_size, seq_len, out_features)
        x = x.view(size_out)

        return x

而Pytorch中FeedForward的實(shí)現(xiàn)(去掉一些細(xì)節(jié))為：

class Linear(Module):
    def __init__(self, in_features: int, out_features: int, bias: bool = True,
                 device=None, dtype=None) -> None:
        super(Linear, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.weight = Parameter(torch.empty((out_features, in_features), **factory_kwargs))
        if bias:
            self.bias = Parameter(torch.empty(out_features, **factory_kwargs))
        else:
            self.register_parameter('bias', None)
        self.reset_parameters()

    def reset_parameters(self) -> None:
        init.kaiming_uniform_(self.weight, a=math.sqrt(5))
        if self.bias is not None:
            fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
            init.uniform_(self.bias, -bound, bound)

    def forward(self, input: Tensor) -> Tensor:
        return F.linear(input, self.weight, self.bias)

我們來看應(yīng)用Conv1D的例子：

embed_dim = 768
conv1d = Conv1D(embed_dim, embed_dim * 3)
# (batch_size, seq_len, embed_dim)
x = torch.rand(2, 5, embed_dim)
# (batch_size, seq_len, embed_dim * 3)
x = conv1d(x)
print(x.shape)

torch.Size([2, 5, 2304])

前饋層

那么就可以應(yīng)用上面的一維卷積來實(shí)現(xiàn)前饋層了：

from torch.nn import functional as F

class MLP(nn.Module):
    def __init__(self, config: GPTConfig) -> None:
        super().__init__()
        embed_dim = config.n_embd
        self.c_fc = Conv1D(embed_dim, embed_dim * 4)
        self.c_proj = Conv1D(embed_dim * 4, embed_dim)
        self.act = F.gelu
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x: Tensor) -> Tensor:
        """

        Args:
            x (Tensor): (batch_size, seq_len, embed_dim)

        Returns:
            Tensor: (batch_size, seq_len, embed_dim)
        """
        # h (batch_size, seq_len, embed_dim * 4)
        h = self.act(self.c_fc(x))
        # h (batch_size, seq_len, embed_dim)
        h = self.c_proj(h)
        return self.dropout(h)

層歸一化

層歸一化這里我們直接使用Pytorch內(nèi)建的torch.nn.LayerNorm。

掩碼多頭注意力

下面我們來實(shí)現(xiàn)掩碼多頭注意力，GPT中的注意力需要防止泄露未來的信息，因此自帶一個(gè)下三角矩陣。

這可以通過以下代碼實(shí)現(xiàn)：

import torch

n_positions = 10

torch.tril(torch.ones(n_positions, n_positions))

tensor([[1., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [1., 1., 0., 0., 0., 0., 0., 0., 0., 0.],
        [1., 1., 1., 0., 0., 0., 0., 0., 0., 0.],
        [1., 1., 1., 1., 0., 0., 0., 0., 0., 0.],
        [1., 1., 1., 1., 1., 0., 0., 0., 0., 0.],
        [1., 1., 1., 1., 1., 1., 0., 0., 0., 0.],
        [1., 1., 1., 1., 1., 1., 1., 0., 0., 0.],
        [1., 1., 1., 1., 1., 1., 1., 1., 0., 0.],
        [1., 1., 1., 1., 1., 1., 1., 1., 1., 0.],
        [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]])

先來看一下初始化方法：

def __init__(self, config: GPTConfig, scale: bool = False) -> None:
    super().__init__()
    self.n_embd = config.n_embd

    assert config.n_embd % config.n_head == 0

    self.scale = scale
    self.n_head = config.n_head

    self.c_attn = Conv1D(self.n_embd, self.n_embd * 3)
    self.c_proj = Conv1D(self.n_embd, self.n_embd)
    # use flash attention or not
    self.flash = hasattr(torch.nn.functional, "scaled_dot_product_attention")
    if not self.flash:
        self.register_buffer(
            "bias",
            torch.tril(torch.ones(config.n_positions, config.n_positions)).view(
                1, 1, config.n_positions, config.n_positions
            ),
            persistent=False,  # will not be saved alongside parameters
        )

    self.attn_dropout = nn.Dropout(config.dropout)
    self.proj_dropout = nn.Dropout(config.dropout)

主要操作是調(diào)用上面實(shí)現(xiàn)的Conv1D，c_attn這樣定義為了可以同時(shí)計(jì)算query,key,value所有的頭，因?yàn)樵贕PT中只有自注意力，由同一個(gè)輸入計(jì)算出不同的query,key,value值，所以可以這樣實(shí)現(xiàn)。

如果還Pytorch2.0及以上的版本，則torch.nn.functional有scaled_dot_product_attention函數(shù)，它利用Flash Attention高效計(jì)算。

否則通過register_buffer將下三角矩陣注冊(cè)為buffer，并且不需要隨著模型參數(shù)保存，轉(zhuǎn)換為(1,1,n_positions,n_positions)的形狀是為了適配批次和多個(gè)頭。

接下來實(shí)現(xiàn)forward()函數(shù)：

 def forward(self, x: Tensor, output_attentions: bool = False) -> list[Tensor]:
        """

        Args:
            x (Tensor): (batch_size, seq_len, n_embd)

        Returns:
            Tensor: (batch_size, seq_len, n_embd) attn_output
            Tensor(optional): (batch_size, n_head, seq_len, seq_len) attn_weights

        """
        # calculate query, key ,value for all heads in batch
        # x (batch_size, seq_len, n_embd * 3)
        x = self.c_attn(x)
        #  query, key, value (batch_size, seq_len, n_embd)
        query, key, value = x.split(self.n_embd, dim=2)
        # query (batch_size,  n_head, seq_len, n_embd / n_head)
        query = self.split_heads(query)
        # key (batch_size, n_head, n_embd / n_head, seq_len)
        key = self.split_heads(key, is_key=True)
        # value (batch_size,  n_head, seq_len, n_embd / n_head)
        value = self.split_heads(value)
        # attn_output (batch_size,  n_head, seq_len, n_embd / n_head)
        attn_outputs = self._attn(query, key, value, output_attentions)
        attn_output = attn_outputs[0]
        # output (batch_size, seq_len, n_embd)
        output = self.merge_heads(attn_output)
        # (batch_size, seq_len, n_embd)
        output = self.c_proj(output)

        output = self.proj_dropout(output)

        outputs = [output] + attn_outputs[1:]
        return outputs

主要過程為：

1. 通過c_attn一次計(jì)算出所有頭的q,k,v值，得到的輸出維度是(batch_size, seq_len, n_embd * 3)；

2. 調(diào)用split在最后一個(gè)維度上將輸出拆分成q,k,v矩陣；

3. 在q,k,v上分別調(diào)用split_heads()拆分成n_head個(gè)頭；

4. 傳入q,k,v調(diào)用_attn()得到注意力計(jì)算結(jié)果；

5. 調(diào)用merge_heads()拼接多頭注意力的結(jié)果；

6. 最后經(jīng)過一個(gè)線性變換c_proj；

split_heads其實(shí)就是一個(gè)變形(view)操作：

    def split_heads(self, x: Tensor, is_key: bool = False) -> Tensor:
        """

        Args:
            x (Tensor): (batch_size, seq_len, n_embd)
            is_key (bool, optional): is key or not. Defaults to False.

        Returns:
            Tensor: (batch_size, n_head, n_embd / n_head, seq_len) if is_key = True ,
              else (batch_size,  n_head, seq_len, n_embd / n_head)
        """
        # (batch_size, seq_len, n_head, n_embd / n_head)
        new_shape = x.size()[:-1] + (self.n_head, x.size(-1) // self.n_head)
        # x (batch_size, seq_len, n_head, n_embd / n_head)
        x = x.view(*new_shape)
        if is_key:
            # (batch_size, n_head, n_embd / n_head, seq_len)
            return x.permute(0, 2, 3, 1)
        # (batch_size,  n_head, seq_len, n_embd / n_head)
        return x.permute(0, 2, 1, 3)

接著就是核心的注意力操作_attn：

def _attn(
    self,
    q: Tensor,
    k: Tensor,
    v: Tensor,
    attention_mask: Tensor = None,
    output_attentions: bool = False,
) -> list[Tensor]:
    """

    Args:
        q (Tensor): (batch_size,  n_head, seq_len, n_embd / n_head)
        k (Tensor): (batch_size, n_head, n_embd / n_head, seq_len)
        v (Tensor): (batch_size,  n_head, seq_len, n_embd / n_head)

    Returns:
        Tensor: (batch_size,  n_head, seq_len, n_embd / n_head) attn_output
        Tensor(optional): (batch_size, n_head, seq_len, seq_len) attn_weights

    """
    if self.flash:
        # 使用flash attention
        attn_output = torch.nn.functional.scaled_dot_product_attention(
            q,
            k,
            v,
            attn_mask=None,
            dropout_p=self.attn_dropout.p if self.training else 0,
            is_causal=True, # 傳入True的話attn_mask必須為None
        )
        weights = None
    else:
        # scores (batch_size,  n_head, seq_len, seq_len)
        scores = torch.matmul(q, k)
        if self.scale:
            scores = scores / math.sqrt(v.size(-1))

        # scores = scores.masked_fill(
        #    self.bias[:, :, : scores.size(-2), : scores.size(-1)] == 0, float("-inf")
        # )
        bias = self.bias[:, :, : scores.size(-2), : scores.size(-1)]
        # more efficient than masked_fill
        scores = scores * bias + -1e9 * (1 - bias)

        # weights (batch_size,  n_head, seq_len, seq_len)
        weights = self.attn_dropout(F.softmax(scores, dim=-1))

        if attention_mask is not None:
            weights = weights + attention_mask

        del scores
        # attn_output (batch_size,  n_head, seq_len, n_embd / n_head)
        attn_output = torch.matmul(weights, v)

    outputs = [attn_output]
    if output_attentions:
        outputs.append(weights)

    return outputs

與上篇文章Transformer中實(shí)現(xiàn)的注意力計(jì)算幾乎沒有變化，對(duì)注意力得分scores進(jìn)行一個(gè)下三角掩碼，這里實(shí)現(xiàn)的時(shí)候采用比masked_fill更高效的乘法和加法的方式。

然后調(diào)用softmax得到注意力權(quán)重，與v矩陣相乘得到最后的注意力輸出。

接下來通過merge_heads拼接多個(gè)注意力頭的結(jié)果：

    def merge_heads(self, x: Tensor) -> Tensor:
        """

        Args:
            x (Tensor):  (batch_size,  n_head, seq_len, n_embd / n_head)

        Returns:
            Tensor: (batch_size, seq_len, n_embd)
        """
        # x (batch_size,  seq_len, n_head, n_embd / n_head)
        x = x.permute(0, 2, 1, 3).contiguous()
        # (batch_size, seq_len, n_embd)
        new_shape = x.size()[:-2] + (x.size(-2) * x.size(-1),)
        return x.view(*new_shape)

其實(shí)也是變形操作。最后經(jīng)過一次線性投影。

此時(shí)模型還未進(jìn)行過非線性操作，為了增強(qiáng)表達(dá)能力，通過前饋層引入非線性操作。

實(shí)現(xiàn)Block

class Block(nn.Module):
    def __init__(self, config: GPTConfig, scale: bool = False) -> None:
        super().__init__()
        n_embd = config.n_embd
        self.attn = Attention(config, scale)
        self.ln_1 = nn.LayerNorm(n_embd)
        self.mlp = MLP(config)
        self.ln_2 = nn.LayerNorm(n_embd)

    def forward(
        self, x: Tensor, attention_mask: Tensor = None, output_attentions: bool = False
    ) -> Tensor:
        """_summary_

        Args:
            x (Tensor): (batch_size, seq_len, n_embd)
            attention_mask (Tensor, optional)
            output_attentions (bool, optional)

        Returns:
            Tensor: (batch_size, seq_len, n_embd) block output
            Tensor(optional): (batch_size, n_head, seq_len, seq_len) attn_weights

        """

        attn_outputs = self.attn(x, attention_mask, output_attentions)
        # a : attention output (batch_size, n_head, seq_len, n_embd / n_head)
        a = attn_outputs[0]

        # resident connection and layer norm
        # n (batch_size, seq_len, n_embd)
        n = self.ln_1(x + a)
        # m (batch_size, seq_len, n_embd)
        m = self.mlp(n)
        # resident connection and layer norm
        # h (batch_size, seq_len, n_embd)
        h = self.ln_2(n + m)

        outputs = [h] + attn_outputs[1:]

        return outputs

Block的實(shí)現(xiàn)就很簡(jiǎn)單，按照架構(gòu)圖實(shí)現(xiàn)即可。這里的attention_mask是用于對(duì)對(duì)填充Token進(jìn)行掩碼。

實(shí)現(xiàn)GPT模型

首先我們要繼承transformers的PreTrainedModel，最終可以將訓(xùn)練好的模型上傳到HuggingFace的Hub上分享給大家。

在這之前我們需要編寫自定義配置，包含構(gòu)建模型所需的所有信息：

from transformers import PretrainedConfig


class GPTConfig(PretrainedConfig):
    model_type = "openai-gpt" # 這個(gè)就是openai的gpt1

    def __init__(
        self,
        vocab_size=5000,
        n_positions=1024,
        n_embd=768,
        n_layer=12,
        n_head=12,
        dropout=0.1,
        initializer_range=0.02,
        **kwargs
    ) -> None:
        """

        Args:
            vocab_size (int, optional): vocabulary size. Defaults to 5000.
            n_positions (int, optional): the maximum sequence length that this model might ever be used with. Defaults to 1024.
            n_embd (int, optional): dimensionality of the embeddings and hidden states. Defaults to 768.
            n_layer (int, optional): number of hidden layers. Defaults to 12.
            n_head (int, optional): number of attention heads for each attention layer. Defaults to 12.
            dropout (float, optional): the dropout probability. Defaults to 0.1.
            initializer_range (tuple, optional): the standard deviation of the truncated_normal_initializer for initializing all weight matrices. Defaults to (0.02,).
        """
        self.vocab_size = vocab_size
        self.n_positions = n_positions
        self.n_embd = n_embd
        self.n_layer = n_layer
        self.n_head = n_head
        self.dropout = dropout
        self.initializer_range = initializer_range

        super().__init__(**kwargs)

編寫自定義配置需要注意三點(diǎn)：

• 繼承自PretrainedConfig；
• __init__方法中必須存在接收任何參數(shù)的kwargs；
• 這些kwargs需要傳遞給父類的__init__方法；

通過繼承我們可以獲得Transformers庫的額外功能，另外兩個(gè)條件是接收PretrainedConfig額外的字段。

有了配置后，我們繼續(xù)編寫GPT模型，同樣繼承類似的PreTrainedModel。先定義一個(gè)基類，主要傳入配置文件、定義參數(shù)初始化方法。

class GPTPreTrainedModel(PreTrainedModel):
    """
    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
    models.
    """

    config_class = GPTConfig
    base_model_prefix = "transformer"

    def __init__(self, config: PretrainedConfig):
        super().__init__(config)

    def _init_weights(self, module):
        if isinstance(module, (nn.Linear, Conv1D)):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)

現(xiàn)在就可以定義我們的GPT模型了：

class GPTModel(GPTPreTrainedModel):

    def __init__(self, config: GPTConfig) -> None:
        super().__init__(config)
        self.config = config
        self.tokens_embed = nn.Embedding(config.vocab_size, config.n_embd)
        self.positions_embed = nn.Embedding(config.n_positions, config.n_embd)

        self.dropout = nn.Dropout(config.dropout)
        self.h = nn.ModuleList(
            [Block(config, scale=True) for _ in range(config.n_layer)]
        )

        self.register_buffer(
            "position_ids", torch.arange(config.n_positions), persistent=False
        )
        self.post_init()

繼承自上面定義的GPTPreTrainedModel，接收配置類。這里負(fù)責(zé)定義詞嵌入和位置編碼，對(duì)于這個(gè)可學(xué)習(xí)的位置編碼，還需要定義表示位置的序列，從0到最大位置，即position_ids。

然后堆疊多層Block，最后調(diào)用self.post_init()，這是PreTrainedModel中為我們實(shí)現(xiàn)的一個(gè)方法，它實(shí)際會(huì)調(diào)用我們自己定義的_init_weights。

再來看前向傳播方法：

def forward(
    self,
    input_ids: torch.LongTensor,
    attention_mask: Tensor = None,
    output_attentions: bool = False,
    output_hidden_states: bool = False,
    return_dict: bool = False,
) -> Union[Tuple[torch.Tensor], BaseModelOutput]:
    """
    Args:
        input_ids (torch.LongTensor): (batch_size, seq_len)
        output_attentions (bool, optional): whether or not to return the attentions tensors of all attention layers. Defaults to False.
        output_hidden_states (bool, optional): whether or not to return the hidden states of all layers. Defaults to False.
        return_dict (bool, optional): whether or not to return a ModelOutput instead of a plain tuple. Defaults to False.



    Returns:
        Union[Tuple[torch.Tensor], BaseModelOutput]: tuple or BaseModelOutput
    """

    input_shape = input_ids.size()

    inputs_embeds = self.tokens_embed(input_ids)
    # generate position ids
    position_ids = self.position_ids[None, : input_shape[-1]]

    position_embeds = self.positions_embed(position_ids)

    hidden_states = inputs_embeds + position_embeds

    hidden_states = self.dropout(hidden_states)

    all_attentions = () if output_attentions else None
    all_hidden_states = () if output_hidden_states else None

    for _, block in enumerate(self.h):
        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)
        outputs = block(hidden_states, attention_mask, output_attentions)
        hidden_states = outputs[0]
        if output_attentions:
            all_attentions = all_attentions + (outputs[1],)

    # add last layer
    if output_hidden_states:
        all_hidden_states = all_hidden_states + (hidden_states,)

    if not return_dict:
        return tuple(
            v
            for v in [hidden_states, all_hidden_states, all_attentions]
            if v is not None
        )

    return BaseModelOutput(
        last_hidden_state=hidden_states,
        hidden_states=all_hidden_states,
        attentions=all_attentions,
    )

這樣我們的GPT編碼層就實(shí)現(xiàn)好了。文章來源地址http://www.zghlxwxcb.cn/news/detail-809066.html

到了這里，關(guān)于手寫GPT實(shí)現(xiàn)小說生成(一)的文章就介紹完了。如果您還想了解更多內(nèi)容，請(qǐng)?jiān)谟疑辖撬阉鱐OY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關(guān)文章，希望大家以后多多支持TOY模板網(wǎng)！