注意力机制在大语言模型中的原理与实现总结

1. 章节介绍

在大语言模型的学习中，理解注意力机制至关重要。本章节旨在深入剖析注意力机制的原理及其在大语言模型中的应用，为构建和优化大语言模型提供理论与实践基础。通过回顾神经网络基础及传统架构的局限性，引出注意力机制的核心内容。在实际应用中，大语言模型广泛用于文本生成、机器翻译、问答系统等场景，而注意力机制是提升这些应用性能的关键因素。无论是程序员在开发相关应用时，还是架构师设计模型架构，都需要对注意力机制有深入理解，因此这部分内容在面试中也备受关注。

核心知识点	面试频率
注意力机制原理	高
与循环神经网络对比	中
注意力机制算法细节	高
注意力机制代码实现	高
注意力机制应用模式	中

2. 知识点详解

在这里插入图片描述

注意力机制原理

传统架构的局限：在传统循环神经网络用于文本翻译的编码器 - 解码器架构里，编码器处理输入文本（如中文）后，将最后一个隐藏状态传给解码器以完成翻译（如翻译成英文）。但随着文本长度增加，期望用一个定长的隐藏状态去包含所有文本知识变得不合理。从数学公式角度看，循环神经网络隐藏状态可分解为词源记忆的加权平均，权重与数据及当前输出位置相关，当距离（i - k）变大时，权重项迅速趋近于0，导致远距离词源对最后隐藏状态贡献的信息量极少，这严重阻碍了模型效果的提升。
注意力机制的革新：注意力机制从宏观上改变了这种状况，它将编码器所有的隐藏状态整合起来，经过一系列处理后再交给解码器。具体实现方式为，针对当前输入，分别计算对每一个隐藏状态的权重，然后基于这些权重对所有隐藏状态进行加权平均，从而得到一个背景向量c。这个背景向量c会作为额外输入提供给解码器，就好像当前输入在背景文本中寻找应该重点关注的词，并依据注意力分布情况整合背景信息。

https://transformers.run/c1/attention/

与循环神经网络对比

循环神经网络的短板：循环神经网络在处理序列数据时，隐藏状态计算存在明显不足。其权重项无法有效表示词元之间的相关关系，并且随着序列长度的增加，远距离词源的信息在传递过程中逐渐丢失，这使得模型难以捕捉到长距离依赖关系。同时，循环神经网络通常以串行方式进行计算，计算效率较低。
注意力机制的优势：注意力机制在计算背景向量时，摆脱了对输出和输入之间距离关系的依赖。其权重项通过特定计算（如key和query的计算结果）能够很好地反映词元之间的相关关系。更为重要的是，注意力机制支持并行计算背景向量，大大提高了计算效率。例如，在处理大规模文本数据时，循环神经网络可能需要花费大量时间按顺序处理每个时间步，而注意力机制可以同时对多个位置的信息进行处理，显著缩短计算时间。

注意力机制算法细节

初始版本：
- 三步计算流程：
  - 计算对齐分数：对于给定的输入i，通过特定函数f（在实际实现中，f通常为内积运算，因为内积在数学上能很好地衡量两个张量之间的相似程度）计算所有隐藏状态h的对齐分数。例如，假设有输入向量i和隐藏状态向量h，通过i和h的内积运算得到一个标量值，该值表示它们之间的对齐程度。
  - 生成权重向量：将计算得到的一系列对齐分数组成一个向量，然后对这个向量应用Softmax函数，得到权重向量WI。Softmax函数的特性使得WI的每一个元素都大于等于0，并且所有元素之和等于1，这使得WI可以作为良好的加权项。
  - 计算背景向量：使用得到的WI对隐藏状态h进行加权平均，从而得到背景向量c。例如，假设有多个隐藏状态h1, h2, h3…，以及对应的权重w1, w2, w3…，背景向量c = w1 * h1 + w2 * h2 + w3 * h3 +…。这个过程中，对齐分数类似于分类模型中计算的逻辑值，它反映了当前输入与背景中每个元素的相似程度评分。
- 存在的问题：
  - 计算资源分配问题：在整个计算过程中，隐藏状态h参与了两处关键计算，一处是在对齐分数的计算，另一处是在计算背景向量时。这种重复使用可能导致h在复杂计算场景下资源分配不足，从而影响模型效果。
  - 架构依赖问题：初始版本的注意力机制紧密依附于编码器和解码器架构，这带来了两个主要问题。其一，限制了注意力机制的使用范围，使其难以独立应用于其他场景；其二，得到的背景向量c在功能上与原本模型中的隐藏变量h存在一定重合，导致整个模型结构变得复杂，且各组件分工不够清晰。
改进版本：
- 独立化设计：后续发展中，将注意力机制从特定架构中独立出来成为一个单独的模型组件，效果更佳。此时，对于一个文本，引入了三个关键张量，分别是value（v）、key（k）、query（q）。
- 新的计算流程：
  - 计算对齐分数：利用k和q来计算对齐分数。这种设计考虑到同一个词源在背景文本和作为输入时所表达的含义有所不同，所以分别用k表示其在背景文本中的含义，用q表示其作为输入时的含义。
  - 生成背景向量：基于计算得到的对齐分数，对v进行加权求和，从而得到背景向量。通过这样的设计，有效解决了短期记忆问题，因为在计算背景向量时不再受输出与输入距离关系的制约。与循环神经网络的隐藏状态计算对比，循环神经网络隐藏状态的权重项无法体现词元之间的相关性，而注意力机制中k和q计算得到的权重项能够很好地反映这种相关性。这里的v类似于词源在不考虑文本背景时所携带的信息量，但具体数学公式与之前有所不同。
- 注意力机制分类：
  - 交叉注意力：主要应用于序列到序列模式。在这种模式下，k和v来自背景文本，而q来自另外一个文本。例如在机器翻译中，源语言文本经过编码器处理后的隐藏状态可作为k和v，目标语言当前位置的输入经处理后作为q。
  - 双向注意力：此时k、v、q都来自同一个文本。在计算对齐分数时，不考虑输入在文本中的位置，即所有位置之间的对齐分数都会正常计算。这意味着在做预测时，可以充分利用输入之前和输入之后的所有文本信息，所以它对应的是自编码模式。比如在文本自动摘要任务中，模型可以利用整个文本的信息来生成摘要。
  - 单向注意力（单向自注意力）：对应自回归模式，k、v、q同样来自同一个文本。但在计算对齐分数时，对于当前输入之后的那些位置，其对齐分数会被强制设为0。例如在文本生成任务中，模型在生成当前词时，只能利用当前位置之前的文本信息，以保证生成过程的因果性。由于要实现GPT模型，后续重点关注单向自注意力机制。

注意力机制代码实现

计算对齐分数：假设k和q形状为BTH（其中b表示批量数据中元素的个数，t表示文本的长度，h表示特征的个数）。根据数学理论，计算对齐分数（即k和q的内积）等同于矩阵乘法。例如，若k是一个形状为(B, T, H)的张量，q是一个形状也为(B, T, H)的张量，将q进行转置操作后（变为(B, H, T)），再与k进行矩阵乘法，即可得到形状为(B, T, T)的对齐分数张量，其中纵向和横向都表示每一个词元，相应元素表示得到的对齐分数。
下三角矩阵与权重分布：
- 利用Softmax特性：Softmax函数对负无穷有特殊处理特性。通过示例说明，假设有一个普通张量a = [1, 2, 3]，对其应用Softmax函数后得到一个元素均为正数且和为1的张量。但当将其中一个元素改为负无穷时，例如a = [1, 2, -inf]，经过Softmax变换后，负无穷所对应的值变为0。利用这一特性，可将对齐分数矩阵转换为下三角矩阵。
- 具体实现步骤：首先随机生成一个方阵来模拟对齐分数矩阵。然后定义一个下三角矩阵，通过该下三角矩阵将对齐分数矩阵的上半部分全部赋值为负无穷。经过Softmax变换后，得到一个下三角矩阵，且该矩阵中每一个行向量的元素都大于等于0，并且行向量元素之和等于1，这就得到了符合要求的权重分布。例如，假设有一个对齐分数矩阵S，通过下三角矩阵mask将S的上半部分元素S[mask == 0] = -inf，再对处理后的S应用Softmax函数，即torch.softmax(S, dim = 1)，得到权重分布。
Softmax方差敏感性：
- 问题表现：Softmax函数对输入的方差非常敏感。通过实验，生成一个方差为1的随机变量x，对其应用Softmax函数，得到的结果较为正常，每个分量大于0且和为1。但当将x的方差放大1000倍后，得到的Softmax结果过于集中于一点，即在某个元素上等于1，其他元素上等于0，这不是我们期望的结果。
- 解决方法：在计算对齐分数时，需要对计算结果进行归一化处理。具体做法是除以h的平方根。例如，在计算对齐分数的代码实现中，将原本的计算式k @ q.transpose(-2, -1)修改为(k @ q.transpose(-2, -1)) / math.sqrt(k.size(-1))，这样处理后得到的对齐分数矩阵的标准差为1，即方差等于1。
函数定义与实现：
- 参数设置：定义一个函数，其参数包括query、value、dropout（用于防止过拟合）和mask。这里的mask起着关键作用，当mask为空时，实现的是双向注意力机制；当mask是一个下三角矩阵时，实现的是单向自注意力机制。默认mask的值为None。
- 张量形状期望：对于query和value，期望的张量形状是BTH；对于mask，期望的张量形状是(T, T)。输出的张量形状依然是BTH。
- 代码实现细节：首先计算对齐分数，代码实现为scores = (query @ key.transpose(-2, -1)) / math.sqrt(query.size(-1))。然后根据mask对scores进行处理，若mask不为None，则scores = scores.masked_fill(mask == 0, float(‘-inf’))。接着计算权重分布，weights = torch.softmax(scores, dim = -1)，如果设置了dropout，则weights = self.dropout(weights)。最后计算背景向量，output = weights @ value，输出的output形状为BTH。
单项式注意力实现：
- 类的定义：定义一个类叫做masked_attention。在初始化函数中，有两个主要参数，分别是length（表示输入的向量长度）和head_size（表示背景向量的长度，关于head_size名字的由来在后续深入讨论中会明确）。
- 模型组件生成：生成key、query和value，这三个模型组件本质上都是线性模型。在代码实现中，通过nn.Linear来创建，并且特别注意不需要它们的截距项。原因主要有两点：一是在大语言模型中通常会使用残差连接这一模型组件来加速模型训练，这种情况下不需要截距项；二是这三个模型组件与文本嵌入类似，而文本嵌入层对应的线性模型在实现时通常也没有截距项。例如，self.key = nn.Linear(length, head_size, bias = False)，self.query = nn.Linear(length, head_size, bias = False)，self.value = nn.Linear(length, head_size, bias = False)。
- 下三角矩阵定义：下三角矩阵在模型中起着辅助运算的作用，它本身不参与模型参数更新。通过torch.tril(torch.ones(sequence_length, sequence_length))来定义一个大的下三角矩阵，在执行向前传播时，从这个大矩阵中截取所需的掩码矩阵mask。例如，mask = self.mask[:x.size(1), :x.size(1)]，这里x是输入张量。
- 随机失活组件定义：定义随机失活组件self.dropout = nn.Dropout(dropout)，用于防止模型过拟合。
- 向前传播算法实现：在向前传播函数中，首先对输入x进行处理，得到形状为(B, T, H)的张量。然后通过定义好的key、query和value模型组件得到相应的三个张量，即k = self.key(x)，q = self.query(x)，v = self.value(x)，它们的形状均为(B, T, H)。接着从大下三角矩阵中截取所需的掩码矩阵mask。需要注意的是，如果输入的序列长度t大于预先定义的sequence_length，计算过程会报错，这体现了注意力机制在处理序列长度上存在限制。最后调用之前定义的注意力计算函数得到最终结果，output = attention(q, k, v, mask = mask, dropout = self.dropout)，输出的output形状为(B, T, H)。在实际应用中，可以运行这部分代码进行简单测试，生成模型实例和测试数据，检查输出是否符合预期以及计算过程是否报错。为了更严谨的测试，还可以参考第三方开源实现，对比两者的计算结果是否一致。

注意力机制应用模式

除了经典的序列到序列模式外，注意力机制可自然地拓展到自回归和自编码模式。在自回归模式中，以文本生成任务为例，模型在预测每一个词元时，传统方式仅依赖当下步骤对应的隐藏状态（即所有已知隐藏状态的最后一个），但这样会丢失之前更多位置的信息。而引入注意力机制后，可以利用当前输入之前的所有隐藏状态进行预测，从而提升模型性能。在自编码模式中，双向注意力机制能让模型在处理输入时，综合利用输入前后的所有文本信息，例如在文本特征提取任务中，可更全面地捕捉文本的语义特征。

3. 章节总结

本章节从大语言模型背景出发，深入探讨注意力机制。介绍了其产生背景，通过与循环神经网络对比突出优势。详细讲解了算法细节，包括初始和改进版本。在代码实现上，逐步阐述了从计算对齐分数到最终实现单向自注意力机制的过程，还介绍了注意力机制在不同应用模式中的特点。

4. 知识点补充

最佳实践

在实际构建大语言模型时，对于注意力机制的应用可参考以下实践：

数据预处理阶段：对输入文本进行合理分词和编码，确保输入数据适合注意力机制计算。例如，在处理长文本时，可采用滑动窗口等技术将长文本分块处理，再分别应用注意力机制，避免一次性处理过长序列导致内存和计算资源不足。以英文文本为例，常用的分词工具如NLTK（Natural Language Toolkit）或spaCy，能够将连续的文本切分成有意义的单词或子词单元。在编码方面，可使用字节对编码（Byte - Pair Encoding，BPE）等方法将分词后的结果转换为模型能够处理的数字表示。当面对超长文本时，设定一个合适的窗口大小，如512个词元，每次只处理窗口内的文本块，计算其注意力权重和相关特征，然后滑动窗口继续处理下一部分，最后将各部分结果整合起来。这样既能有效利用注意力机制的优势，又能克服长序列带来的计算挑战。
模型搭建阶段：根据任务需求选择合适的注意力机制类型。若为文本翻译等序列到序列任务，可优先考虑交叉注意力机制；若是文本生成等自回归任务，单向自注意力机制更为合适。同时，合理设置多头注意力机制的头数，通过实验确定最优参数，以平衡计算成本和模型性能。在文本翻译任务中，源语言和目标语言之间存在着对应关系，交叉注意力机制能够很好地捕捉这种跨文本的依赖关系，帮助模型在翻译时参考源语言的整体信息。对于文本生成任务，单向自注意力机制符合文本生成的因果性，即生成当前词时只能依赖之前已生成的词。在设置多头注意力机制的头数时，一般从较小的数值如2、4开始尝试，逐步增加头数并观察模型在训练集和验证集上的性能表现，包括损失值、准确率、BLEU（bilingual evaluation understudy）分数等指标。当头数增加到一定程度后，模型性能提升可能不再明显，反而会带来计算量的大幅增加，此时可确定一个较为合适的头数，如8或16。
训练阶段：采用合适的优化算法和超参数调整策略，如使用AdamW优化器，结合学习率衰减策略，防止模型在训练过程中过拟合，提高注意力机制的学习效果。并且，在训练过程中监控注意力权重分布，分析模型对不同位置元素的关注情况，以便及时调整模型结构或训练策略。AdamW优化器是在Adam优化器的基础上改进而来，它能够更好地处理权重衰减问题，防止模型参数在训练过程中过度增长导致过拟合。学习率衰减策略可以在训练初期设置一个较大的学习率，使模型快速收敛，随着训练的进行，逐渐减小学习率，让模型在接近最优解时更加稳定地收敛。例如，可采用指数衰减策略，学习率 = 初始学习率 * 衰减率 ^（当前步数 / 衰减步数）。通过可视化工具（如TensorBoard）监控注意力权重分布，观察模型在不同训练阶段对输入文本中各个位置词元的关注程度。如果发现模型对某些位置或某些类型的词元关注异常，如过度关注高频词而忽略低频词，可尝试调整模型结构，如增加更多的注意力层或调整注意力机制的参数，或者优化训练数据，增加低频词的样本数量等。

编程思想指导

模块化编程思想：在实现注意力机制代码时，将不同功能模块分开编写，如将对齐分数计算、权重生成、背景向量计算等功能分别封装成函数或类方法。这样不仅使代码结构清晰，易于维护和调试，还方便在不同项目中复用这些模块。例如，在实现单向自注意力机制时，将生成下三角矩阵的功能封装成一个独立函数，在其他需要下三角矩阵的场景中也可直接调用。在Python中，可以定义一个函数generate_triangular_mask，输入为序列长度，函数内部通过torch.tril(torch.ones(seq_length, seq_length))生成下三角矩阵并返回。在注意力机制的主函数中，当需要生成掩码矩阵时，直接调用generate_triangular_mask函数即可。对于对齐分数计算，可定义一个compute_attention_scores函数，输入为query和key张量，内部实现矩阵乘法和归一化操作，返回计算得到的对齐分数张量。通过这种模块化设计，当需要修改某一功能的实现细节时，只需在对应的模块中进行调整，而不会影响到整个代码的其他部分。
抽象与泛化思想：从注意力机制的多种应用模式中抽象出通用的计算逻辑，将其设计为可配置参数的通用函数或类。比如，通过传入不同的mask参数，使同一个注意力机制实现函数既能处理单向自注意力，也能处理双向自注意力和交叉注意力。这样在面对不同的任务和数据时，代码具有更好的适应性和扩展性。以实现注意力机制的类为例，在类的初始化函数中接收一个mask_type参数，根据这个参数的值来决定生成何种类型的掩码矩阵。如果mask_type为’unidirectional’，则生成下三角矩阵用于单向自注意力；如果为’bidirectional’，则生成全1矩阵（即不进行掩码操作）用于双向自注意力；如果为’cross’，则根据具体的交叉注意力场景生成相应的掩码矩阵。在类的前向传播函数中，根据生成的掩码矩阵进行后续的注意力计算。这样，通过简单地修改mask_type参数，就可以将同一个类应用于不同的注意力机制场景，大大提高了代码的复用性和灵活性。
性能优化思想：考虑到注意力机制的计算量较大，在编码过程中注重性能优化。例如，合理利用矩阵运算库（如PyTorch或TensorFlow中的矩阵运算函数）进行并行计算，减少循环操作；在计算对齐分数时，通过归一化等操作避免数值不稳定问题，提高计算效率和稳定性。同时，在处理大规模数据时，采用分布式计算或模型并行策略，提升模型训练和推理速度。在PyTorch中，矩阵乘法操作torch.matmul比使用Python循环实现的矩阵乘法快得多，因为它利用了底层的CUDA（Compute Unified Device Architecture）加速库进行并行计算。在计算对齐分数时，除以特征维度的平方根进行归一化，这不仅能保证数值稳定性，还能使Softmax函数的输出分布更加合理。当处理大规模数据时，可采用分布式训练框架，如Horovod，它能够将训练数据分布到多个计算节点上并行处理，大大缩短训练时间。对于模型并行策略，可以将模型的不同层分配到不同的计算设备（如多个GPU）上进行计算，避免单个设备内存不足的问题，同时提高计算效率。例如，将注意力机制层和后续的前馈神经网络层分别放在不同的GPU上进行计算，通过合理的数据传输和同步机制，实现模型的高效运行。

5. 程序员面试题

简单题

请简述注意力机制的基本思想。
答案：注意力机制的基本思想是将编码器所有隐藏状态合起来，通过计算当前输入对每个隐藏状态的权重，加权平均得到背景向量，作为解码器的附加输入，从而解决传统架构中定长隐藏状态难以表示所有信息以及远距离词源信息丢失的问题。

中等难度题

对比循环神经网络和注意力机制在处理序列数据时的优缺点。
答案：
- 循环神经网络：优点是能处理不定长序列，可捕捉序列中的时间依赖关系；缺点是随着序列增长，远距离词源对隐藏状态贡献信息少，存在梯度消失或梯度爆炸问题，且计算通常是串行的，效率较低。
- 注意力机制：优点是计算背景向量时不依赖输入输出距离，权重项能反映词元相关关系，背景向量可并行计算，计算效率高；缺点是处理序列长度有限制，依赖特定架构时可能导致模型结构复杂。
解释多头注意力机制的工作原理及其优势。
答案：
- 工作原理：将输入投影到多个子空间，每个子空间独立进行注意力计算，即分别计算每个子空间的对齐分数、权重和背景向量，最后将各个子空间的结果拼接起来。在具体实现中，通过多个线性变换将输入分别转换为多个不同的query、key和value张量，然后对每个子空间的query、key和value执行常规的注意力计算流程，最后将所有子空间的输出按特定维度拼接。
- 优势：能同时关注输入序列的不同方面，捕捉更丰富的信息，提高模型的表示能力和性能，有助于模型学习到更复杂的模式和关系。不同的头可以关注输入的不同特征，例如在文本处理中，有的头关注语义信息，有的头关注语法结构信息，综合起来使模型对文本的理解更全面。

高难度题

在实现注意力机制代码时，如何优化计算以提高效率和稳定性？请详细说明。
答案：
- 计算优化：
  - 利用矩阵运算库进行并行计算，例如在计算对齐分数时，将向量运算转换为矩阵乘法，充分利用GPU等并行计算设备的优势，提高计算速度。在PyTorch中，使用torch.matmul函数进行矩阵乘法，相比于使用Python循环实现的向量运算，能够极大地加速计算过程，因为torch.matmul会自动调用GPU的并行计算资源。
  - 避免不必要的循环操作，如在生成下三角矩阵和计算权重分布时，尽量使用向量化操作，减少计算开销。例如，生成下三角矩阵时，使用torch.tril(torch.ones(seq_length, seq_length))这种向量化操作，而不是通过循环逐个元素赋值。计算权重分布时，对整个分数矩阵直接应用Softmax函数，而不是循环处理每个元素。
- 稳定性优化：
  - 在计算对齐分数时，除以特征维度h的平方根进行归一化，保证输入Softmax函数的方差为1，防止结果过于集中，提高数值稳定性。这是因为Softmax函数对输入的数值范围较为敏感，归一化可以使输入数值处于一个合理的范围，使得Softmax的输出分布更加稳定和合理。
  - 合理处理数值范围，例如在将对齐分数矩阵上半部分赋值为负无穷时，确保数值在计算机可表示范围内，避免溢出或下溢问题。在实际编程中，使用系统提供的表示负无穷的常量（如Python中float('-inf')），并且在进行大规模矩阵运算时，注意数据类型的选择和数值范围的监控，防止出现数值异常导致计算错误。
设计一个基于注意力机制的文本分类模型，阐述模型架构、关键组件及其作用，并给出简要的代码实现思路（使用Python和PyTorch框架）。
答案：
- 模型架构：模型可由嵌入层、注意力层、全连接层组成。
- 关键组件及其作用：
  - 嵌入层：将文本中的每个词转换为向量表示，捕捉词的语义信息。通过查找预训练的词向量表（如Word2Vec、GloVe等）或在训练过程中学习词向量，将输入的离散词转换为连续的向量表示，使得模型能够处理和理解文本中的语义内容。
  - 注意力层：采用自注意力机制，计算文本中各个词之间的注意力权重，突出重要词汇，使模型能关注到文本中不同位置的关键信息。通过计算每个词与其他词之间的注意力分数，生成权重向量，对文本的词向量表示进行加权求和，从而突出与文本分类任务相关的重要词汇的特征。
  - 全连接层：对注意力层输出进行分类，得到文本分类结果。将注意力层输出的特征向量映射到类别空间，通过一系列的线性变换和激活函数，输出每个类别的预测概率。
- 代码实现思路：

import torch
import torch.nn as nnclass AttentionTextClassifier(nn.Module):def __init__(self, vocab_size, embedding_dim, hidden_dim, num_classes):super(AttentionTextClassifier, self).__init__()self.embedding = nn.Embedding(vocab_size, embedding_dim)self.attention = nn.MultiheadAttention(embedding_dim, num_heads = 1)self.fc = nn.Linear(embedding_dim, num_classes)def forward(self, x):x = self.embedding(x)x = x.permute(1, 0, 2)  # 调整维度顺序以适应注意力层输入attn_output, _ = self.attention(x, x, x)attn_output = attn_output.mean(dim = 0)  # 对注意力输出求平均output = self.fc(attn_output)return output

在上述代码中，vocab_size是词汇表大小，embedding_dim是词向量维度，hidden_dim是隐藏层维度，num_classes是分类类别数。首先通过嵌入层将输入文本转换为向量，然后经过多头自注意力层，对注意力输出求平均后通过全连接层进行分类。

BV1DW421R7rz

代码

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import numpy as np# 自定义数据集类
class TextDataset(Dataset):def __init__(self, texts, labels):# 初始化文本数据和标签self.texts = textsself.labels = labelsdef __len__(self):# 返回数据集的长度return len(self.texts)def __getitem__(self, idx):# 根据索引获取数据集中的文本和标签text = torch.tensor(self.texts[idx], dtype=torch.long)label = torch.tensor(self.labels[idx], dtype=torch.long)return text, label# 基于注意力机制的文本分类模型类
class AttentionTextClassifier(nn.Module):def __init__(self, vocab_size, embedding_dim, hidden_dim, num_classes, num_heads=1):# 调用父类的构造函数super(AttentionTextClassifier, self).__init__()# 嵌入层，将文本中的每个词转换为向量表示self.embedding = nn.Embedding(vocab_size, embedding_dim)# 多头注意力层，计算文本中各个词之间的注意力权重self.attention = nn.MultiheadAttention(embedding_dim, num_heads)# 全连接层，对注意力层输出进行分类self.fc = nn.Linear(embedding_dim, num_classes)def forward(self, x):# 通过嵌入层将输入文本转换为向量x = self.embedding(x)# 调整维度顺序以适应注意力层输入x = x.permute(1, 0, 2)# 经过多头注意力层，得到注意力输出和注意力权重attn_output, _ = self.attention(x, x, x)# 对注意力输出求平均attn_output = attn_output.mean(dim=0)# 通过全连接层进行分类output = self.fc(attn_output)return output# 训练模型的函数
def train_model(model, train_loader, criterion, optimizer, device, epochs):# 将模型设置为训练模式model.train()for epoch in range(epochs):running_loss = 0.0for texts, labels in train_loader:# 将数据移动到指定设备（如 GPU）texts = texts.to(device)labels = labels.to(device)# 清零梯度optimizer.zero_grad()# 前向传播，得到模型的输出outputs = model(texts)# 计算损失loss = criterion(outputs, labels)# 反向传播，计算梯度loss.backward()# 更新模型参数optimizer.step()running_loss += loss.item()# 打印每个 epoch 的损失print(f'Epoch {epoch + 1}/{epochs}, Loss: {running_loss / len(train_loader)}')# 评估模型的函数
def evaluate_model(model, test_loader, device):# 将模型设置为评估模式model.eval()correct = 0total = 0with torch.no_grad():for texts, labels in test_loader:# 将数据移动到指定设备texts = texts.to(device)labels = labels.to(device)# 前向传播，得到模型的输出outputs = model(texts)# 获取预测的类别_, predicted = torch.max(outputs.data, 1)total += labels.size(0)correct += (predicted == labels).sum().item()# 计算准确率accuracy = 100 * correct / totalprint(f'Accuracy: {accuracy}%')# 主函数
if __name__ == "__main__":# 超参数设置vocab_size = 10000embedding_dim = 128hidden_dim = 128num_classes = 2num_heads = 1epochs = 10batch_size = 32learning_rate = 0.001# 模拟生成一些训练数据num_samples = 1000texts = [np.random.randint(0, vocab_size, 10) for _ in range(num_samples)]labels = np.random.randint(0, num_classes, num_samples)# 创建数据集和数据加载器dataset = TextDataset(texts, labels)train_size = int(0.8 * len(dataset))test_size = len(dataset) - train_sizetrain_dataset, test_dataset = torch.utils.data.random_split(dataset, [train_size, test_size])train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)# 检查是否有可用的 GPUdevice = torch.device("cuda" if torch.cuda.is_available() else "cpu")# 创建模型实例，并将其移动到指定设备model = AttentionTextClassifier(vocab_size, embedding_dim, hidden_dim, num_classes, num_heads).to(device)# 定义损失函数和优化器criterion = nn.CrossEntropyLoss()optimizer = optim.Adam(model.parameters(), lr=learning_rate)# 训练模型train_model(model, train_loader, criterion, optimizer, device, epochs)# 评估模型evaluate_model(model, test_loader, device)

代码说明

TextDataset 类：继承自 torch.utils.data.Dataset，用于封装文本数据和对应的标签，方便后续使用 DataLoader 进行批量加载。
AttentionTextClassifier 类：定义了基于注意力机制的文本分类模型，包含嵌入层、多头注意力层和全连接层。
train_model 函数：用于训练模型，在每个 epoch 中进行前向传播、损失计算、反向传播和参数更新。
evaluate_model 函数：用于评估模型的准确率，在测试集上进行预测并计算准确率。
主函数：设置超参数，生成模拟数据，创建数据集和数据加载器，定义模型、损失函数和优化器，然后进行模型训练和评估。