Transformer处理文本分类实例(Pytorch)

项目开源:GIT地址

https://gitee.com/TomCoCo/transformer_test

Transformer处理文本分类实例

Transformer的应用太多了:

1. 文字生成图像(DALL-E,Midjourney)
2. 文本翻译(transformer原文的目的)
3. chartGPT(以给定的文本预测下一个单词)

本文主旨基于介绍Transformer,做一个简单的文本分类任务来说明

参考网站

• LLM可视化模型,描述了chartGPT的模型结构,非常详细

llm可视化模型

• 3Blue1Brown 的Transformer解释的非常棒,推荐一看

【官方双语】直观解释注意力机制，Transformer的核心 | 【深度学习第6章】 https://www.bilibili.com/video/BV1TZ421j7Ke)

我们构建一个实例问题,预测AG_NEWS的文本分类

AG_NEWS数据集介绍

1. 数据格式如下, csv格式,分别为: 分类,标题,描述

“3”,“Wall St. Bears Claw Back Into the Black (Reuters)”,“Reuters - Short-sellers, Wall Street’s dwindling\band of ultra-cynics, are seeing green again.”

2. AG News数据集包含四个主要的类别：1->世界（World）、2->体育（Sports）、3->商业（Business）,4->科技（Sci/Tech）。
3. 每个类别包含30,000个训练样本和1,900个测试样本，总共120,000个训练样本和7,600个测试样本。
4. AG News数据集由AG’s Corpus of News Articles收集整理而来，该语料库由ComeToMyHead（一个自2004年7月起运行的学术新闻搜索引擎）在超过一年的时间内从2000多个新闻源中收集了超过100万篇新闻文章。

预测目标

在7600个数据的验证集,总体acc@1>90%,即判断对分类的概率为90%. (参数量在2000w左右)

总体思路(简述)

• 首先建立一个涵盖所有词的词表,可以理解为一个大词典,初始状态只有所有词的目录,但是没有词的释义,
• 理论上来说,一门完全陌生从未见过的外语,只要给出相当数量的文本,依旧可以破译出每一个词的含义,训练的过程,就是计算机破译词语(token)的过程.
• 而推理的的过程,即翻阅已经破译好之后的词典,然后组合上下文确定具体语义的过程
• 例如下面段(来自于dict.baidu.com):初始情况下,我们只建立了一个序号->词 的映射关系,释义都是空的.

• 我们观察这个词的含义,足足有5个;含义复杂的词汇甚至有十几种意义完全不同的解释.
• 例如"意思":

• 训练的第一个目标,就是将所有的词的语义,进行数值映射. 即使用n维(词嵌入特征维度)的tensor,用以存储语义信息.
• 从ids->到n维tensor的过程就称之为词嵌入,初始情况下,目标tensor称之为词嵌入张量,又称之为词嵌入矩阵/词嵌入向量(单个词),初始为随机值(无意义信息,只是占位).
• 词嵌入的维度,称之为词嵌入(特征)维度决定着对于任意一个词,语义丰富程度的上限,如果维度过低,无法容纳过多的语义信息
• 训练的过程,就是从每一个语句中获取到每一个词的含义,然后更新词嵌入张量.词嵌入向量的更新是为了使其在所有语境中表现得更好，而不是完全捕捉每个具体上下文的所有信息。词嵌入捕捉的是词的通用语义特征，而上下文中的具体语义会通过Transformer的自注意力机制来捕捉。
• 训练阶段:同一个词在不同的语句中的含义会不同,故而词嵌入张量会被不断更新,以叠加更多的语义在其中(散布在不同的层参数)
• !!注意!! 词嵌入层,保存的是一个token_id到通用语义的映射
• 推理阶段,使用id计算词嵌入结果, 这里的处理过程是使用嵌入层计算初始词嵌入张量(注意,这里由于被训练过,故而是有值的,即词典已经构造完成),然后使用自注意力层和mlp,更新词向量的数据,使之更加匹配当前的上下文语义.(即可以理解为,通过上下文选择组合一些释义,作为这个词当前的释义)
• 总的来说,使用嵌入层将token_id转为具有通用含义的tensor,使用自注意力机制和后续的网络层会调整这些向量,使其更好地反映当前输入的上下文语义。这样就获取到了当前句子(准确来说是上下文长度),每一个词的在当前语境下实际含义.
• 有了一句话的每个词的含义,做一个句子长度(上下文长度)->分类的全连接即可完成分类工作

主要流程

总体类似于下图,图像来自于https://bbycroft.net/llm,点击nano-gpt模型,本文主要结构和nano-gpt模型的形状一致

数据预处理

dataset构建(不是重点)

我们简单描述dataset,以及datasetloader的构建,说白了就是把数据封装为标签和文本的过程,以下是batch_size=2的示例

Label: Size([2])

tensor([2, 1])

Texes:tulpe: Size:2(batch_size=2)

(

‘Fosters counts cost of poor wine market Brewing giant Fosters today said that higher annual profits from sales of beer overseas had been eroded by a sharp drop in US wine trade.’,

‘Carter - joined the Jets. (Getty Images) Less than three weeks after being released by the Dallas Cowboys, quarterback Quincy Carter has landed with the New York Jets. Carter arrived in New York on Tuesday and signed a one-year contract.’

)

from torchtext import datasets
# 可以直接使用torchtext构建
train_data = datasets.AG_NEWS(root=train_data_path, split='train')
valid_data = datasets.AG_NEWS(root=train_data_path, split='test')
self.train_data_loader = DataLoader(dataset=train_data, batch_size=self.batch_size, shuffle=True)
self.valid_data_loader = DataLoader(dataset=valid_data, batch_size=self.batch_size, shuffle=True)
self.train_data_size = 120000
self.valid_data_size = 7600

构建词表

1. 首先,我们想处理任意的文本输入,需要告诉模型,你需要认识哪些词汇(即token),也就是给模型一份(词)字典,

2. 分词的主要过程:

   首先将一句话切分成若干的语义单元(token),例如下面这句话的分词方式

   我/买了/一个/金色/的/黄铜/材质/的/东方明珠/塔/模型(分词方式1)

   我/买了/一个/金色的/黄铜材质的/东方明珠塔/模型(分词方式2)

   • 分词方式: 每一个词,有其本身的含义在,合适的分词是处理的第一步,

     例如:方式2就会比方式1,构建出更厚的一本词表,但是相对的,每个词的含义更加明确不会混淆;

     极端的,如果针对每一个字,单独切分,构建出更薄的一本词表也是可行的.就会从词典变成字典. 每个字的含义会更加混淆不清楚;

     这样带来的影响,就是模型需要通过更多的上下文,去推断这个字的含义,因为单字的含义大多数时候是模糊不清的.

     你看新华字典可比现代汉语词典要小的多了

   • 词表大小: 再看这句话,其中包含了英文和中文,如果目标是兼容2种文字,则词表会建立的更大,相当于建立了一本<<现代汉语和英语词典>>

     熟悉/markdown/语法/可以/让/文章/排版/更/便捷/美观

     确定你的需要处理的文字的范围,以确定词表大小,把所有的词(token)都从0-length编个有序列表,这样一个空的词表就建好了. 相当于建立了一本词典,但是目前只有目录,所有词条的含义尚未定义.这些词条定义,需要在学习中逐渐定义出来

   • 分词器: 中文分词器jieba,英文分词器BertTokenizerFast.本文针对AG_news处理,故而使用BertTokenizerFast

   • 特殊词处理:我构建词表时,最简单的方式就是将训练集的所有句子,统一执行分词,然后去重,就很容易得到一个训练集上出现的所有单词列表,这个是没有问题的. 但是预测推理的过程中,无法保证新的词,一定在词表中出现. 就如现在的很多网络新词,你查<<现代汉语词典>>也是查不到的,所幸,如果训练集足够大,那么词表足够大,缺少的词也不会很多,故而影响不是很大.

     • 未知词标记 [UNK]
     • 填充词标记 [PAD] 和上下文长度有关,后面会解释

3. 为处理AG_NEWS, 由于是全英文的,此处我们不去手动构建词表,使用BertTokenizerFast提供的词表,词表长度:30522; 使用以下方法可以将文本数组,使用 分词->映射->填充->转换格式 的处理流程处理为了tensor格式的下标.

   转换示例:

   tulpe: Size:2(batch_size=2)

   (

   ‘Fosters counts cost of poor wine market Brewing giant Fosters today said that higher annual profits from sales of beer overseas had been eroded by a sharp drop in US wine trade.’,

   ‘Carter - joined the Jets. (Getty Images) Less than three weeks after being released by the Dallas Cowboys, quarterback Quincy Carter has landed with the New York Jets. Carter arrived in New York on Tuesday and signed a one-year contract.’

   )

   ->Tensor:Size([2, 512])

   tensor([[ 101, 6469, 2015, …, 0, 0, 0],
   [ 101, 5708, 1011, …, 0, 0, 0]])

   转换代码:

   from transformers import BertTokenizerFast
   # 分词器和词表,注意这里的词表,仅仅使用了一个分词功能和id映射功能,没有使用到词嵌入的映射哦
   tokenizer = BertTokenizerFast.from_pretrained('bert-base-uncased')
   # 分词,max_length是上下文长度
   encoded_batch = tokenizer.batch_encode_plus(
                   texts,
                   max_length=512,
                   truncation=True,
                   padding='max_length',
                   return_tensors='pt'
               )
   # 转化为了tensor格式的下标
   input_ids = encoded_batch['input_ids'].to(self.device)
   

4. 上下文长度:

• 如果我们把一个自然句,认为是一次的训练样本,那么有以下推论(训练和推理都存在)
• 如果一句话非常长,每一个词受到的影响就会越多,因为上下文越多,但是消耗计算资源也会变多
• 如果一句话非常短,每一个词受到的影响就会越少,因为上下文越少,但是消耗计算资源也会变少
• 即:上下文长度会影响到一个词会影响到/受影响多远距离的词.
• 实际情况下,句子有长有短,当前模型无法处理一个变长的语句,故而将超出的词截去,过短的使用填充词填充(上文提到的 [PAD] ,下标为0)
• !!非常重要!! 我们规定一个合适的长度,将所有输入的文本都处理成等长的tensor即可,称之为上下文长度(序列长度)

编写处理模型

执行词嵌入

• 实际处理的是token_id到通用语义的映射这么一个过程
• 训练阶段,从无通用语义,逐渐使用已有的上下文去推断语义是什么
• 推理阶段,使用token_id,提取通用语义.
• 故而形状可以轻松推断,即词表宽度x词嵌入深度
• 词表宽度由构建词表的方式决定,词嵌入深度决定词的语义丰富程度的上限
• 在pytorch中定义嵌入层非常简单

vocab_size = self.tokenizer.vocab_size
# 我这里设计深度为300,参数选择可以见附录参数参考
embedding_feature_dim = 300
self.embedding = nn.Embedding(vocab_size, embedding_feature_dim)

• 此时,数据格式变化为

->Tensor:Size([2, 512])

tensor([[ 101, 6469, 2015, …, 0, 0, 0],
[ 101, 5708, 1011, …, 0, 0, 0]])

->Tensor:Size([2, 512, 300])

tensor([[[-0.0749, -1.5324, 0.0588, …, -1.8509, 0.4253, -1.5777],
[ 0.5036, 1.1159, -1.0334, …, 0.7873, -1.9076, -0.5892],
[ 1.6551, -1.7568, 1.0481, …, -1.8378, -1.4752, -1.7096],
…,
[-1.0518, 0.9318, -0.3377, …, 0.6822, -1.4751, 0.9430],
[-1.0518, 0.9318, -0.3377, …, 0.6822, -1.4751, 0.9430],
[-1.0518, 0.9318, -0.3377, …, 0.6822, -1.4751, 0.9430]],
[[-0.0749, -1.5324, 0.0588, …, -1.8509, 0.4253, -1.5777],
[ 0.7489, -0.8386, 0.9028, …, 3.1699, 0.8188, -0.2222],
[-0.0861, 0.1495, -0.1039, …, -0.1902, -2.5695, 0.0744],
…,
[-1.0518, 0.9318, -0.3377, …, 0.6822, -1.4751, 0.9430],
[-1.0518, 0.9318, -0.3377, …, 0.6822, -1.4751, 0.9430],
[-1.0518, 0.9318, -0.3377, …, 0.6822, -1.4751, 0.9430]]],
grad_fn=)

• 注意,这里非常重要:我们获得了一个batch_size x 上下文长度 x 词特征嵌入深度的tensor(词嵌入结果),接下来的操作,本质都是在更新它的过程,它就代表这这句话的语义. 此处的示例形状为Size([2, 512, 300])

位置编码(PositionalEncoding)

• 这段512长度的句子,在300的特征维度中,并没有记录位置信息,即包含谁在前,谁在后的信息. 因为是从词嵌入层获取的
• 故而使用编码技术,执行位置编码,将位置信息写入到特征维度中,此时张量形状不变,而且无可学习参数,只是单纯的更新数据,以嵌入位置信息
• 模块代码如下:

class PositionalEncoding(nn.Module):
    """
    位置编码器
    """

    def __init__(self, d_model, max_len=512, device=None):
        super(PositionalEncoding, self).__init__()
        self.encoding = torch.zeros(max_len, d_model)
        self.register_buffer('positional_encoding', self.encoding)

        position = torch.arange(0, max_len).unsqueeze(1).float()
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model))
        self.encoding[:, 0::2] = torch.sin(position * div_term)
        self.encoding[:, 1::2] = torch.cos(position * div_term)
        self.encoding = self.encoding.unsqueeze(0)  # Add batch dimension
        self.encoding = self.encoding.to(device)

    def forward(self, x):
        return x + self.encoding[:, :x.size(1)]

(*核心)多层Transformer模块

• 我们先来看一下这个模型的形状:其中有一个6x,代表以下的模块重复建立6次.(互为独立参数),重复次数为注意力模块的深度 (depth)

(transformer): Transformer(
(layers): ModuleList(
(0-5): 6 x ModuleList(
(0): AttentionMulti(
(w_qkv): Linear(in_features=300, out_features=3072, bias=True)
(attention_to_embedding): Linear(in_features=1024, out_features=300, bias=True)
(dropout): Dropout(p=0.0, inplace=False)
)
(1): Sequential(
(0): LayerNorm((300,), eps=1e-05, elementwise_affine=True)
(1): Linear(in_features=300, out_features=1200, bias=True)
(2): GELU(approximate=‘none’)
(3): Dropout(p=0.1, inplace=False)
(4): Linear(in_features=1200, out_features=300, bias=True)
(5): Dropout(p=0.1, inplace=False)
)
)
)

• 我们先看单个的Transformer模块,为一个多头注意力模块和MLP层构成

(0): AttentionMulti(
(w_qkv): Linear(in_features=300, out_features=3072, bias=True)
(attention_to_embedding): Linear(in_features=1024, out_features=300, bias=True)
(dropout): Dropout(p=0.0, inplace=False)
)
(1): Sequential(
(0): LayerNorm((300,), eps=1e-05, elementwise_affine=True)
(1): Linear(in_features=300, out_features=1200, bias=True)
(2): GELU(approximate=‘none’)
(3): Dropout(p=0.1, inplace=False)
(4): Linear(in_features=1200, out_features=300, bias=True)
(5): Dropout(p=0.1, inplace=False)
)

多头自注意力模块

• 先看AttentionMulti,结构貌似非常简单,忽略掉dropout,就是2个全连接层而已,观察2个全连接层

  • w_qkv 从词嵌入维度(3000)到3072维度
  • attention_to_embedding 从1024维度到300维度
  • 我们剖开w_qkv,其实是3个全连接层,为了计算方便组合在一起计算罢了,等效于

  wq = nn.Linear(300, 1024)
  wk = nn.Linear(300, 1024)
  wv = nn.Linear(300, 1024)
  

  • 实际的此层的操作,可以猜测,貌似是使用2个全连接层,升维再降维: 即从300维->1024维,然后又从1024维度->300维

• 好的,我们先不看这个多层多头的Transformer模块,我们先去理解,这个模块想要做什么

  1. 核心任务改变词嵌入张量,以适应上下文的语义. 核心在改变一词,我们需要获取到这个变化量dt,可以看出,变化量的维度必须和词嵌入张量维度一致,不然无法叠加上去

  2. 这个模块的功能就是从上下文中提取到合适的语义,然后计算这个变化量

  3. 如何计算呢,首先我们需要衡量,一句话中,那一个词会对当前词产生影响,影响因子(QK)有多少,对于塔这个词来说,影响可能是这样的

     比如: 我/买了/一个/黄金的/埃菲尔/塔.

     对于词 “塔” 来说

     上下文中: 一个,黄金的,埃菲尔 都对塔产生了明显语义影响,且影响很大,而: 我,买了 则影响相对较小.

     影响源	影响目标	attention
     我	塔	0.09
     买了	塔	0.01
     一个	塔	0.2
     黄金的	塔	0.2
     埃菲尔	塔	0.5
     塔	塔	0

     上表中的是单个对"塔"单个的影响,而且是归一化的概率值,下面是一句话中的所有词构成的注意力矩阵,logits值

  4. 其次,我们需要知道,每一个词对目标词的**影响量(V)**是多少,例如""黄金的"对"塔"的影响,就应该使之具有一种金灿灿,值钱,黄色等等语义的赋予.注意这个影响量也是具有维度的,维度越高,那么可以赋予的语义空间就越大.

  5. 所有的词的**影响量(V)**x,影响因子(QK),对塔有影响的数据.

  6. 原始数据切分为token,token选择对应的词向量表示(从嵌入矩阵中),此时的向量只编码了基础单词含义,是没有上下文影响存在的

     上下文长度,每一个词能够看到的上下文对他的影响距离

     上下文长度的token,输入到模型中,和解嵌入矩阵相乘,得到预测单词分布

     输入6,深度128,即输入6x128记为X
     X流经 注意力模块 还是X ,但是被编码了上下文信息,即数据改变了,即一句话里面的每一个词,含有了其上下文的语义,而不是一个单独的词存在
     埃菲尔铁塔模型 -> 拆分为 埃菲尔/铁塔/模型 铁塔(初始值), 埃菲尔对其影响, 模型对其影响

     生成Q的矩阵, ,词向量为128x1时,为了生成Q,则

     生成Q,即有哪些特征会影响到本词

     生成K,即本词具有哪些特征

     QK的乘积,即本词和哪些其他的词匹配上了,匹配度是多少,即本词需要的特征与其他词的提供特征对应上了.

     比如: 我/买了/一个/黄金的/埃菲尔/塔

     上下文中明显: 一个,黄金的,埃菲尔 都对塔产生了语义影响,且影响很大,而: 我,买了 则影响相对较小.

     换一个角度来说,

     对于词: 塔 来说,他的Q(128维度),第一维是期待获取形状,第二维度材质,第三维度颜色,第四维度数量,第五维度位置等等等等

     对于词:黄金的 来说,他K(128维度)第一维描述为重量,第二维度描述为颜色,第三维描述为形容词等等等等(注意,这里是特征名,不是特征值)

     对于词:黄金的 来说,他V(128维度),则是他的K的对应的每一个维度的特征值,例如第一维是:很重,第二维是:金色,第三维:是

     执行点积,即执行两者的相似度度量,如果相似,则匹配较好,故而这些词对他的影响会更大,

  7. 注意力矩阵大小等于上下文长度的平方

  8. 具体的计算方案,

     1. 构建$W_{QKV}$三个全连接层,输入参数即词嵌入矩, 用以计算$Q,K,V$

        $W_Q$矩阵 :生成Q矩阵 :词嵌入特征维度->注意力特征维度

        $W_K$矩阵 :生成K矩阵 :词嵌入特征维度->注意力特征维度

        $W_V$矩阵 :生成V矩阵 :词嵌入特征维度->注意力特征维度

     2. 计算注意力矩阵

        $$Attention(Q,K)=softmax(\frac{Q{K}^T}{\sqrt{d_k}})$$

        其中$d_k$为注意力特征维度(此例为128)

        除以 $\sqrt{d_k}$ 可以使点积的结果保持在一个较为稳定的范围内，确保 softmax 函数的输出不至于变得过于极端。

     3. 注意力矩阵乘以V

        $$Attention(Q,K,V)=softmax(\frac{Q{K}^T}{\sqrt{d_k}})V$$

        得出的值即应该添加到原数据上的变化量

     4. 注意,此时的维度不一致,需要将变化量结果从: 注意力特征维度->词嵌入特征维度

        wq = nn.Linear(300, 128)
        wk = nn.Linear(300, 128)
        wv = nn.Linear(300, 128)
        
        attention_to_embedding = nn.Linear(128, 300)
        
        ......省略若干
        
        # 使用WK,WQ 300->128 , 将嵌入矩阵6x512x300分别全连接为6x512x128,
        # 512来自于上下文长度,128来自于kq特征维度(attn_feature_dim).
        # 计算KQV结果: 6x512x128 ,
        q_tensor = wq(token_tensor)
        k_tensor = wk(token_tensor)
        v_tensor = wv(token_tensor)
        
        # 点积qk(注意需要转置矩阵),得注意力矩阵(attention matrix),6x512x512. 即计算kq的128维度相关性
        # 点积的结果值,和attn_feature_dim相关,维度越高,则点积的结果越大
        attention_tensor = torch.matmul(q_tensor, k_tensor.transpose(-1, -2))
        # 为平衡点积结果,故而除以根号下attn_feature_dim  6x512x512->6x512x512
        attention_tensor = attention_tensor / torch.sqrt(torch.tensor(128, dtype=torch.float32))
        # 对结果做softmax得到概率值(6x512x512)
        attention_tensor = torch.softmax(attention_tensor, dim=-1)
        # 做叉乘,得6x512x128,即在这个注意力网络中,对每一个的词向量的改变值.
        v_output = torch.matmul(attention_tensor, v_tensor)
        # 注意这个变化量形状(6x512x128)和词向量形状(6x512x300)不一致. 再做一个线性变换,从注意力特征维度转为300词嵌入维度
        attention_delta = attention_to_embedding(v_output)
        # 将变化量添加到原数据上
        token_tensor.add(attention_delta)
        

• 以上是单头单层注意力模块的内容,数据的变化如下

6x512x300 + (6x512x300 -> 6x512x128 -> 6x512x300)

• 多(n)头,即将$QKV$复制多n份,例如8头,则 存在8组$W_Q,W_K,W_V$,则输出结构x8,又由于kqv结构一致,故而合并为一个tensor处理则之前的问题可以解释:
• (w_qkv): Linear(in_features=300, out_features=3072, bias=True) 其中的3072应该分解为 128x8x3 : 128维度的注意力模块深度,8头,qkv三组参数
• (attention_to_embedding): Linear(in_features=1024, out_features=300, bias=True) 单头的计算应该是6x512x128 -> 6x512x300,即128->300维度,由于是8头,则是8x128 = 1024输入层

MLP模块

• MLP看上去貌似更简单,实际也是更简单,其实就是2个全连接顺序执行,先升维度,激活函数,再降维

  Sequential(
  (0): LayerNorm((300,), eps=1e-05, elementwise_affine=True)
  (1): Linear(in_features=300, out_features=1200, bias=True)
  (2): GELU(approximate=‘none’)
  (3): Dropout(p=0.1, inplace=False)
  (4): Linear(in_features=1200, out_features=300, bias=True)
  (5): Dropout(p=0.1, inplace=False)
  )

• 首先目的:

  • 引入语法语义层次的语义信息到句子中,而不是简单词义的叠加
  • 这里引入的GELU的激活函数,引入了非线性映射,
  • 将不同的特征加权重新映射,MLP可以学习到不同特征之间的相关性和相互作用，使模型能够更好地理解和表示数据的复杂结构。
  • 例如，在语言模型中，不同的词语嵌入向量在MLP中经过变换后，可以捕捉到词语之间的语法和语义关系。

分类层

1. 经过若干次的transformer模块,句子的合适的语义信息,已经被捕捉到了,保存在6x512x300的tensor中,我们处理的是一个4分类任务,即我们需要一个6x4的可能性tensor
2. 执行分类操作,直接执行 300->4的分类是缺少意义的,是针对一句话的每一个词(512)分别求分类,需要在第1维度作平均,平均一句话的内容,即6x512x300->6x4
3. 这样再执行一次全连接 300->4即可

# 执行分类操作,直接执行 300->4的分类是缺少意义的,是针对一句话的每一个词(512)分别求分类,需要在第1维度作平均,平均一句话的内容
x = x.mean(dim=1)
# 全连接执行分类操作,得分类的logtis值
return self.fc(x)

附录,完整代码

import math
import time
from datetime import datetime

import torch
from functorch.einops import rearrange

from torch import nn, optim
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from torchtext import datasets
from transformers import BertTokenizerFast


class PositionalEncoding(nn.Module):
    """
    位置编码器
    """

    def __init__(self, d_model, max_len=512, device=None):
        super(PositionalEncoding, self).__init__()
        self.encoding = torch.zeros(max_len, d_model)
        self.register_buffer('positional_encoding', self.encoding)

        position = torch.arange(0, max_len).unsqueeze(1).float()
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model))
        self.encoding[:, 0::2] = torch.sin(position * div_term)
        self.encoding[:, 1::2] = torch.cos(position * div_term)
        self.encoding = self.encoding.unsqueeze(0)  # Add batch dimension
        self.encoding = self.encoding.to(device)

    def forward(self, x):
        return x + self.encoding[:, :x.size(1)]


class AttentionMulti(nn.Module):
    def __init__(self, embedding_feature_dim, attention_feature_dim, heads, dropout=0.):
        """
        :param embedding_feature_dim: 词嵌入特征维度: 例300
        :param attention_feature_dim: 注意力特征维度(每一个头),即dk 例:128
        :param heads: 多头注意力的头数 例:8
        """
        super(AttentionMulti, self).__init__()
        # qkv合一,形状一致,故而深度为attention_feature_dim * heads * 3
        self.w_qkv = nn.Linear(embedding_feature_dim, attention_feature_dim * heads * 3)
        self.heads = heads
        self.dk = attention_feature_dim / heads
        self.attention_to_embedding = nn.Linear(attention_feature_dim * heads, embedding_feature_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask):
        # 6x512x300 -> 6x512x(128x8x3)
        qkv = self.w_qkv(x)
        # 沿着qkv切开 3(tuple)x6x512x(128x8)
        qkv = qkv.chunk(3, dim=-1)
        # 分离出heads提到前面来 3(tuple)x6x8x512x128
        # b (batch_size); h (heads)头数 ;n (sequence length)序列长度(上下文长度);d (dimension) 注意力特征维度
        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=self.heads), qkv)

        # 6x8x512x128 叉乘 6x8x512x128 ->6x8x512x512,得128维组合的相关性
        attention_tensor = torch.matmul(q, k.transpose(-1, -2))
        attention_tensor = attention_tensor / torch.sqrt(torch.tensor(self.dk, dtype=torch.float32))

        # 在计算注意力权重时，将填充部分的权重设置为负无穷大，以使 softmax 计算后变为零。从而使填充不会影响到其他数据
        if mask is not None:
            # 使用填充掩码调整注意力权重
            mask = mask.unsqueeze(1).unsqueeze(2)  # Shape: [b, 1, 1, n]
            attention_tensor = attention_tensor.masked_fill(mask == 0, float('-inf'))

        # 对结果做softmax得到概率值(6x8x512x512)
        attention_tensor = attention_tensor.softmax(dim=-1)
        attention_tensor = self.dropout(attention_tensor)

        # 做叉乘,得6x8x512x128,即在这个注意力网络中,对每一个的词向量的改变值.
        out = torch.matmul(attention_tensor, v)
        # 将多头获取到的变化合并到最后的改变特征值维度,6x8x512x128 -> 6x512x(128x8)
        # 即每个头都为原始的词嵌入提供了128维度的变化量,总计8个头,故而提供128x8,所以合并这些变化量
        out = rearrange(out, 'b h n d -> b n (h d)', h=self.heads)
        out = self.attention_to_embedding(out)
        out = self.dropout(out)
        # 输出尺寸和输入一致 6x512x300
        return out


class Transformer(nn.Module):
    def __init__(self, depth, embedding_feature_dim, attention_feature_dim, heads, mlp_feature_dim, dropout=0.):
        """
        :param depth: 模型深度:例如3
        :param embedding_feature_dim: 词嵌入特征维度: 例300
        :param attention_feature_dim: 注意力特征维度(每一个头),即dk 例:128
        :param heads: 多头注意力的头数 例:8
        :param mlp_feature_dim: mlp的特征维度:例1200
        """
        super(Transformer, self).__init__()
        self.layers = nn.ModuleList([])
        for _ in range(depth):
            # 多头注意力
            attention = AttentionMulti(embedding_feature_dim, attention_feature_dim, heads)
            # 每个MLP一个ln,2个全连接,一个激活函数构成
            mlp = nn.Sequential(
                nn.LayerNorm(embedding_feature_dim),
                nn.Linear(embedding_feature_dim, mlp_feature_dim),
                nn.GELU(),
                nn.Dropout(dropout),
                nn.Linear(mlp_feature_dim, embedding_feature_dim),
                nn.Dropout(dropout),
            )
            self.layers.append(nn.ModuleList([attention, mlp]))
        # 深度次数后的的Transformer,加上一个层归一化
        self.norm = nn.LayerNorm(embedding_feature_dim)

    def forward(self, x, mask):
        for attention, mlp in self.layers:
            x = x + attention(x, mask)
            x = x + mlp(x)
        return self.norm(x)


class TextClassifyTransformer(nn.Module):
    def __init__(self, class_num, vocab_size, sequence_length, depth, embedding_feature_dim, attention_feature_dim,
                 heads, mlp_feature_dim, dropout=0., emb_dropout=0., device=None):
        """
        :param class_num: 分类数量
        :param vocab_size: 词表大小,
        :param sequence_length: 序列长度(上下文长度)
        :param depth: Transformer模型深度:例如3,即自注意力模块和mlp重复多少次
        :param embedding_feature_dim: 词嵌入特征维度: 例300
        :param attention_feature_dim: 注意力特征维度(每一个头),即dk 例:128
        :param heads: 多头注意力的头数 例:8
        :param mlp_feature_dim: mlp的特征维度:例1200
        :param dropout: 丢弃率.这里一共会使用4处,1.注意力矩阵计算完成,对注意力矩阵执行. 2.从注意力转为词向量的全连接层之后 3.mlp网络第一次全连接,激活后,4.mlp第二次全连接后
        :param emb_dropout: 对抗过拟合,此处为编码完成之后的dropout,即对词嵌入后执行的
        """
        super(TextClassifyTransformer, self).__init__()
        # 建立词嵌入层,词表长度,嵌入维度
        self.embedding = nn.Embedding(vocab_size, embedding_feature_dim)
        # 归一化
        self.ln1 = nn.LayerNorm(embedding_feature_dim)
        # 位置编码器
        self.pos_encoder = PositionalEncoding(d_model=embedding_feature_dim, max_len=sequence_length, device=device)
        self.emb_dropout = nn.Dropout(emb_dropout)

        # Transformer模块
        self.transformer = Transformer(depth, embedding_feature_dim, attention_feature_dim, heads, mlp_feature_dim,
                                       dropout)

        self.fc = nn.Linear(embedding_feature_dim, class_num)

    def forward(self, x, mask):
        # 执行词嵌入,将输入input转对应词向量 6x512x300
        x = self.embedding(x)

        # 嵌入位置  6x512x300
        x = self.pos_encoder(x)

        x = self.emb_dropout(x)

        # 执行layerNorm,即 6x512x300 的300维度执行归一化,即对一句话的每一个词向量执行归一化,2组权重参数β和γ
        x = self.ln1(x)

        # 执行Transformer模块,得 6x512x300
        x = self.transformer(x, mask)

        # 执行分类操作,直接执行 300->4的分类是缺少意义的,是针对一句话的每一个词(512)分别求分类,需要在第1维度作平均,平均一句话的内容
        x = x.mean(dim=1)

        # 全连接执行分类操作,得分类的logtis值
        return self.fc(x)


class TextClassify:
    def __init__(self, train_data_path=None, workers=8,
                 batch_size=128, epochs_num=2, lr=1e-5, target_path="./target"):
        """
        :param train_data_path: 数据集的路径
        :param workers: 加载器工作线程数
        :param batch_size:
        :param epochs_num:
        :param lr:
        :param target_path: 权重文件保存位置
        """
        # 定义设备
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        print(f"use divice:{self.device}")
        # tensorboard 记录器
        self.writer = SummaryWriter(log_dir='./log/' + time.strftime('%m-%d_%H.%M', time.localtime()))

        # 定义目标目录
        self.target_path = target_path

        # 设定超参数:minibatch大小,迭代次数,学习率,正则惩罚
        self.batch_size = batch_size
        self.epochs_num = epochs_num
        self.lr = lr
        self.workers = workers

        # 分词器和词表,注意这里的词表,仅仅使用了一个分词功能和id映射功能,没有使用到词嵌入的映射哦
        self.tokenizer = BertTokenizerFast.from_pretrained('bert-base-uncased')
        # 词表长度:30522
        print(f"The vocabulary size is: {self.tokenizer.vocab_size}")

        # 构造模型
        self.model = TextClassifyTransformer(4, self.tokenizer.vocab_size, sequence_length=512, depth=6,
                                             embedding_feature_dim=300, attention_feature_dim=128, heads=8,
                                             mlp_feature_dim=1200, dropout=0.1, emb_dropout=0.1, device=self.device)

        self.model.to(self.device)

        # 损失函数,优化器
        self.criterion = nn.CrossEntropyLoss()
        self.optimizer = optim.AdamW(self.model.parameters(), lr=self.lr, weight_decay=1e-2)
        self.lr_scheduler = optim.lr_scheduler.StepLR(self.optimizer, step_size=1, gamma=0.7)

        # 数据集,预处理方式
        self.train_data_loader = None
        self.valid_data_loader = None
        self.train_data_size = None
        self.valid_data_size = None

        # 临时变量,用以计数,训练和验证执行的batch的index,在一次训练中始终从0开始,一直递增
        self.valid_batch_index = 0
        self.train_batch_index = 0

        # 加载数据
        self.__load_data(train_data_path)

    def __load_data(self, train_data_path):
        train_data = datasets.AG_NEWS(root=train_data_path, split='train')
        valid_data = datasets.AG_NEWS(root=train_data_path, split='test')
        self.train_data_loader = DataLoader(dataset=train_data, batch_size=self.batch_size, shuffle=True)
        self.valid_data_loader = DataLoader(dataset=valid_data, batch_size=self.batch_size, shuffle=True)
        # self.train_data_size = len(train_data)
        # self.valid_data_size = len(valid_data)
        self.train_data_size = 120000
        self.valid_data_size = 7600

    def train_model(self):
        print("start train model...")
        print(f"hyper-parameters: batch_size:{self.batch_size}; "
              f"epochs_num:{self.epochs_num}; lr:{self.lr};")

        # 批次最佳
        best_epoch_acc_rate_valid = 0

        # 重置批次计数
        self.train_batch_index = 0
        self.valid_batch_index = 0

        for epoch in range(self.epochs_num):
            epoch_loss_train, epoch_acc_rate_train = self.__do_train(epoch)
            # 更新优化器
            self.lr_scheduler.step()
            # 执行验证
            epoch_loss_valid, epoch_acc_rate_valid = self.__do_valid(epoch)

            # 记录每个epoch的优化情况
            self.writer.add_scalars("epoch_loss", {"train": epoch_loss_train}, epoch)
            self.writer.add_scalars("epoch_loss", {"valid": epoch_loss_valid}, epoch)
            self.writer.add_scalars("epoch_acc", {"train": epoch_acc_rate_train}, epoch)
            self.writer.add_scalars("epoch_acc", {"valid": epoch_acc_rate_valid}, epoch)

            print(f"epoch {epoch}/{self.epochs_num - 1} : "
                  f"epoch_loss_train:{epoch_loss_train:.4f}; epoch_acc_rate_train:{epoch_acc_rate_train:.4f}; "
                  f"epoch_loss_valid:{epoch_loss_valid:.4f}; epoch_acc_rate_valid:{epoch_acc_rate_valid:.4f}; ")

            # 更新最优参数
            if epoch_acc_rate_valid >= best_epoch_acc_rate_valid:
                best_epoch_acc_rate_valid = epoch_acc_rate_valid
                torch.save(self.model.state_dict(),
                           f"{self.target_path}/state_dict_{datetime.now().strftime('%Y%m%d%H%M%S')}"
                           f"_{epoch_acc_rate_train:.4f}_{best_epoch_acc_rate_valid:.4f}.pth")

    def __do_train(self, epoch):
        """
        训练数据
        :param epoch:
        :return:
        """
        # 进入训练模式,计算梯度
        self.model.train()

        epoch_loss_sum = 0
        epoch_acc_sum = 0

        # minibatch计数
        ix = 0
        for label, texts in self.train_data_loader:
            label = label - 1
            label = label.to(self.device)

            # 分词,找词表对应下标,还有填充,这里一个方法一起做了
            encoded_batch = self.tokenizer.batch_encode_plus(
                texts,
                max_length=512,
                truncation=True,
                padding='max_length',
                return_tensors='pt'
            )
            # token的id,6x512
            input_ids = encoded_batch['input_ids'].to(self.device)
            attention_mask = encoded_batch['attention_mask'].to(self.device)

            with (torch.set_grad_enabled(True)):
                # 传入填充过的mask,计算分类
                output = self.model(input_ids, attention_mask)

                loss = self.criterion(output, label)
                _, prediction = torch.max(output, 1)
                # 执行反向传播
                loss.backward()
                self.optimizer.step()
                self.optimizer.zero_grad()

                # 计算当前小批量的正确率
                batch_acc = prediction.eq(label).sum()
                # 计算批次准确率
                # 当前批量大小,由于数据集可能都不足一个batch_size,所以要获取到准确的batch_size
                now_batch_size = label.shape[0]

                # 将损失乘上size是为了在计算平均损失时，考虑到每个样本对损失的贡献。
                epoch_loss_sum += (loss * now_batch_size)
                epoch_acc_sum += batch_acc
                # 记录事件
                self.writer.add_scalars("batch_train", {"loss": loss}, self.train_batch_index)
                self.writer.add_scalars("batch_train", {"acc": batch_acc / now_batch_size}, self.train_batch_index)

                print('[%d/%d][%d/%d]\t%s\t loss: %.4f\t acc_rate: %.4f'
                      % (epoch, self.epochs_num, ix, self.train_data_size, datetime.now(), loss.item(), batch_acc / now_batch_size))

                self.train_batch_index += 1
                ix += self.batch_size

        epoch_loss = epoch_loss_sum / self.train_data_size
        epoch_acc_rate = epoch_acc_sum / self.train_data_size
        return epoch_loss, epoch_acc_rate

    def __do_valid(self, epoch):
        """
        验证过程
        :param epoch:
        :return:
        """
        # 进入推理模式,不计算梯度
        self.model.eval()

        epoch_loss_sum = 0
        epoch_acc_sum = 0

        ix = 0
        for label, texts in self.valid_data_loader:
            label = label - 1
            label = label.to(self.device)

            # 分词,找词表对应下标,还有填充,这里一个方法一起做了
            encoded_batch = self.tokenizer.batch_encode_plus(
                texts,
                max_length=512,
                truncation=True,
                padding='max_length',
                return_tensors='pt'
            )
            # token的id,6x512
            input_ids = encoded_batch['input_ids'].to(self.device)
            attention_mask = encoded_batch['attention_mask'].to(self.device)

            with (torch.set_grad_enabled(False)):
                # 传入填充过的mask,计算分类
                output = self.model(input_ids, attention_mask)

                loss = self.criterion(output, label)
                _, prediction = torch.max(output, 1)
                # 计算当前小批量的正确率
                batch_acc = prediction.eq(label).sum()
                # 计算批次准确率
                # 当前批量大小,由于数据集可能都不足一个batch_size,所以要获取到准确的batch_size
                now_batch_size = label.shape[0]

                # 将损失乘上size是为了在计算平均损失时，考虑到每个样本对损失的贡献。
                epoch_loss_sum += (loss * now_batch_size)
                epoch_acc_sum += batch_acc
                # 记录事件
                self.writer.add_scalars("batch_valid", {"loss": loss}, self.valid_batch_index)
                self.writer.add_scalars("batch_valid", {"acc": batch_acc / now_batch_size}, self.valid_batch_index)

                print('[%d/%d][%d/%d]\t%s\t loss: %.4f'
                      % (epoch, self.epochs_num, ix, self.valid_data_size, datetime.now(), loss.item()))

                self.valid_batch_index += self.batch_size
                ix += 1

        epoch_loss = epoch_loss_sum / self.valid_data_size
        epoch_acc_rate = epoch_acc_sum / self.valid_data_size
        return epoch_loss, epoch_acc_rate


"""
多头示例
"""
if __name__ == '__main__':
    cc = TextClassify(train_data_path="./resources", batch_size=2, epochs_num=2, lr=1e-4, target_path="./target")
    cc.train_model()

附录: 指明padding的影响

指明padding的index是什么,会加速收敛速度. 不指明模型其实也可以习得,

附录:Transformer 关键参数选择参考

对于构建 Transformer 模型的关键参数选择，以下是建议的设定：

1. 词嵌入维度 (embedding_feature_dim)

• 推荐值: 256, 300, 或 512
• 选择依据:
  • 256: 适用于资源有限或基础任务。
  • 300: 通常的选择，平衡性能和资源消耗。
  • 512: 高性能需求的复杂任务。

2. 多头注意力的头数 (heads)

• 推荐值: 4, 8, 或 12
• 选择依据:
  • 4: 计算开销较小，适合基础任务。
  • 8: 常见选择，适合大部分 NLP 任务。
  • 12: 适用于需要更高表达能力的复杂任务。

3. 注意力模块的深度 (depth)

• 推荐值: 4, 6, 或 12
• 选择依据:
  • 4: 适合轻量级模型，任务较简单。
  • 6: 通用选择，平衡性能和复杂性。
  • 12: 高性能需求的复杂任务，如机器翻译、问答系统。

4. 注意力特征维度 (attention_feature_dim)

• 推荐值: 64, 128, 或 256
• 选择依据:
  • 64: 适合资源有限或模型简化的场景。
  • 128: 常见选择，适合大部分 NLP 任务。
  • 256: 高性能需求的复杂任务。

5. MLP 特征维度 (mlp_feature_dim)

• 推荐值: 2 * embedding_feature_dim 到 4 * embedding_feature_dim
• 选择依据:
  • 2 * embedding_feature_dim: 适合资源有限或模型简化的场景。
  • 4 * embedding_feature_dim: 提供更高的模型表达能力，适合复杂任务。

示例配置

• 中等规模模型:
  • embedding_feature_dim = 300
  • heads = 8
  • depth = 6
  • attention_feature_dim = 128
  • mlp_feature_dim = 600 (或 900)

这些推荐值可以根据具体任务、硬件资源和期望的模型性能进行调整。初始选择后，可以通过实验和调优进一步优化。

附录,训练结果

1. 深度为6和8影响不大
2. 示例代码为深度8时训练2轮,acc@1大约91%,耗时一个上午(H100)