自注意力机制（Self-Attention）

自注意力机制（Self-Attention）是一种在序列或空间维度上衡量不同位置之间关联的计算方法，其核心思想是每个位置的表示不仅由其自身的前向信息决定，还由其他所有位置的表示影响。

自然语言处理上的应用

作用功能

自注意机制的作用：能使模型在每个时间步上考虑整个序列的信息，使得信息在模型内部能够做到长距离传播，更为充分地考虑上下文中的依赖关系，将上文中的所有相关理解都融入现有正在思考的词中。
例如模型在推理侦探小说时，考虑最后一句“凶手是”中的token“是”时，会考虑所输入序列中的所有相关信息，使得token“是”的嵌入向量包含远超自身原有含义的内容，也就是在模型运行的高维向量空间中更具有指向性，有“注意力”。

更具体的功能如下：
信息聚合：自注意力机制能够在输入序列中不同位置之间计算关联性，聚合重要的信息。这种方式使得模型在生成某个元素时，可以参考整个输入序列中的相关信息，而不仅仅是局部上下文。

动态权重分配：自注意力通过计算不同输入元素之间的相似度，为它们分配不同的权重。这样，模型可以根据上下文动态调整注意力焦点，强调重要信息，抑制不相关的信息。

并行计算：与传统的循环神经网络（RNN）相比，自注意力机制能够支持更高效的并行计算，因为它不依赖于输入的顺序处理。这使得处理长序列时更具效率。

长距离依赖：自注意力机制有效解决了序列模型中长距离依赖的问题。在处理长文本时，RNN往往难以捕捉远距离的信息，而自注意力直接连接了序列中所有位置，使得模型能够轻松捕获长距离的依赖关系。

通用性：自注意力机制不仅在NLP中表现出色，如在Transformer模型中得到了广泛应用，同时也被越来越多地应用于计算机视觉等领域，例如Vision Transformer（ViT）等。

计算过程

将计算过程分为六部分：

1. 输入表示：
   词嵌入：将输入序列中的每个词转换为一个固定维度的向量。如下图中为"Thinking" “Machines”
   位置编码：为了保留序列中词的位置信息，通常会添加位置编码。
2. 查询、键和值的生成：
   查询向量Q：通过输入表示与查询权重矩阵相乘得到。与权重矩阵相乘这一过程可以使用神经网络中的线性层来表示。
   键向量K：通过输入表示与键权重矩阵相乘得到。
   值向量V：通过输入表示与值权重矩阵相乘得到。

图片来源：【超详细】【原理篇&实战篇】一文读懂Transformer
3. 注意力分数计算：
点积：计算查询向量Q和键向量K的点积，得到注意力分数。
缩放：为了避免点积结果过大，通常会除以键向量维度的平方根。
掩码（可选）：在某些情况下（如在Transformer模型中的掩码多头自注意力），会使用掩码来忽略某些位置的注意力分数。
4. 注意力权重计算：
Softmax函数：对缩放后的注意力分数应用Softmax函数，得到归一化的注意力权重。
5. 加权求和：
加权值求和：将注意力权重与值向量相乘，然后对所有位置求和，得到最终的输出表示。
6. 输出表示：
经过自注意力机制处理后的输出表示，可以用于后续的层或任务。

最终输出如下图自注意力头：

多头注意力机制（Multi-Head Attention）

多头注意力机制（Multi-Head Attention）是自注意力机制的一个扩展，它允许模型在不同的表示子空间中并行地学习信息。这种机制可以捕捉到不同类型的依赖关系，并增加模型的表达能力。

实际计算过程只需重复上文单个注意力头z的计算过程,拼接之后与权重矩阵W相乘得到信息矩阵Z。
权重矩阵W在模型中联合训练，它与模型的其他参数一起通过反向传播算法进行更新和优化。


下面为使用PyTorch实现多头自注意力机制的简单代码示例

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

# 定义超参数
batch_size = 2  # 批量大小
seq_length = 5  # 序列长度
d_model = 64     # 特征维度
n_heads = 8      # 注意力头数量
d_k = d_model // n_heads  # 每个头的维度

# 创建自定义词向量，设定一些简单的关系
# 也可使用randn直接生成随机张量，因为这一过程很难从计算结果上做到可视化，所以确实没必要像此处定义关系
# input_embeddings = torch.randn(batch_size, seq_length, d_model)
input_embeddings = torch.tensor([
    [
        [1.0] + [0.0] * 63,  # 第一个词
        [0.0, 1.0] + [0.0] * 62,  # 第二个词
        [0.0, 0.0, 1.0] + [0.0] * 61,  # 第三个词
        [0.0, 0.0, 0.0, 1.0] + [0.0] * 60,  # 第四个词
        [0.0, 0.0, 0.0, 0.0, 1.0] + [0.0] * 59  # 第五个词
    ],
    [
        [0.5, 0.5] + [0.0] * 62,  # 第一个词
        [0.0, 0.0, 0.5, 0.5] + [0.0] * 60,  # 第二个词
        [0.0, 0.0, 0.0, 0.5, 0.5] + [0.0] * 59,  # 第三个词
        [0.0, 0.0, 0.0, 0.0, 0.0, 0.5, 0.5] + [0.0] * 57,  # 第四个词
        [0.0] * 64  # 第五个词
    ]
], dtype=torch.float32)  # 确保数据类型为float32

# 定义线性变换层 使用线性层表示向量与对应权重矩阵相乘的过程
query_layer = nn.Linear(d_model, d_model)  # 查询的线性变换
key_layer = nn.Linear(d_model, d_model)    # 键的线性变换
value_layer = nn.Linear(d_model, d_model)  # 值的线性变换

# 定义前向传播函数
def forward(input_embeddings):
    # 通过线性变换生成查询、键和值
    Q = query_layer(input_embeddings)  # 查询向量
    K = key_layer(input_embeddings)      # 键向量
    V = value_layer(input_embeddings)    # 值向量

    # 转换形状以进行多头注意力计算
    Q = Q.view(batch_size, seq_length, n_heads, d_k).transpose(1, 2)  # (batch_size, n_heads, seq_length, d_k)
    K = K.view(batch_size, seq_length, n_heads, d_k).transpose(1, 2)  # (batch_size, n_heads, seq_length, d_k)
    V = V.view(batch_size, seq_length, n_heads, d_k).transpose(1, 2)  # (batch_size, n_heads, seq_length, d_k)

    # 计算注意力分数，得分矩阵形状为 (batch_size, n_heads, seq_length, seq_length)
    attention_scores = Q @ K.transpose(-2, -1)  # 矩阵乘法
    attention_scores = attention_scores / (d_k ** 0.5)  # 应用缩放

    # 计算注意力权重
    attention_weights = F.softmax(attention_scores, dim=-1)  # 应用 softmax

    # 应用注意力权重到值上
    output = attention_weights @ V  # (batch_size, n_heads, seq_length, d_k)

    # 将多头输出合并回一个张量
    output = output.transpose(1, 2).contiguous().view(batch_size, seq_length, d_model)  # (batch_size, seq_length, d_model)

    return Q, K, V, attention_weights, output

# 运行前向传播
Q, K, V, attention_weights, output = forward(input_embeddings)

# 输出结果
print("输入词嵌入：\n", input_embeddings)
print("\n查询向量 Q：\n", Q)
print("\n键向量 K：\n", K)
print("\n值向量 V：\n", V)
print("\n注意力权重：\n", attention_weights)
print("\n输出结果：\n", output)

图像识别上的应用

Vision Transformer

Vision Transformer (ViT) 是一种用于图像分类的深度学习模型，它将 Transformer 架构（最初用于自然语言处理）应用于计算机视觉任务。ViT 通过将图像划分为多个小块（patches）并将它们视为序列数据，借助自注意力机制来处理图像特征。

计算过程

1. 模型架构
   输入图像处理（Patch Embedding）：
   首先，将输入图像 (224 × 224 )划分为多个小块（patches），例如 (16× 16) 像素的小块。每个小块被展平并线性映射为一个固定维度的向量，一个patch[16,16,3]铺平后为[768] 维。这些处理称为“Patch Embedding”。
   特征展平(Flatten):
   将所有小块（以及类标识符）连接在一起，形成一个序列。序列中的小块数量取决于图像中的小块总数（197：196个小块 + 1个类标识符）。接着，对其进行展平（Flatten）。
   类标识符（Class Token)：
   ViT 还引入一个特殊的“类标识符”（Class Token），它被连接到所有小块的序列前面。这个类标识符将用于最终的分类任务。
   位置编码（Position Embedding）：
   为了将空间位置信息引入到模型中，ViT 添加了位置编码（Position Embedding）。这些编码与每个小块的特征相加，以使模型可以理解保留小块在图像中的位置信息。

2. 特征提取和编码
   标准化（Layer Norm）：
   层归一化（Layer Norm）。
   Transformer 编码器（transformer Encoder）：
   ViT 的核心是多个 Transformer 编码器（encoder blocks），通常由 (12) 个堆叠的编码器组成。每个编码器内部包含两个主要组件：多头自注意力机制（Multi-Head Attention）和前馈神经网络（MLP Block）。
   编码器通过自注意力机制捕捉小块之间的全局上下文关系，通过前馈神经网络进行进一步特征变换，最后经过层归一化。
   

Dropout 机制：
在多头自注意力和 MLP Block 之后，使用 Dropout 层来防止过拟合。
3. 特征输出
提取类标识符特征（Extract Class Token）：
通过所有编码器处理后的特征序列中，提取类标识符对应的特征。这个特征包含了整个输入图像的全局信息。
多层感知机头（MLP Head）：
最后，提取的类标识符特征将进入多层感知机头（MLP Head），通常由一个全连接层（Linear）和激活函数（tanh）构成。在经过最终的线性变换后，输出图像的分类结果。
4. 训练与性能
ViT 在大规模图像数据集（如 ImageNet）上进行训练，在充分的数据量下能表现出与传统卷积神经网络（CNN）相媲美甚至更优的性能。此外，由于其模型结构的特性，ViT 更适合于大规模训练。

应用特点与示例代码

5. 优点与缺点
   优点：
   能够捕捉全局上下文信息，实现对图像的有效理解。
   模型模块化设计，使其容易扩展和修改。
   在大规模计算上表现良好，借助于 Transformer 的并行化能力。
   缺点：
   对于小规模数据集可能表现不佳，需要大量标注数据才能发挥优势。
   计算复杂度较高，尤其是在输入序列很长的情况下，可能导致更高的计算资源需求。
6. 应用场景
   ViT 不仅适用于图像分类任务，还可以扩展到目标检测、图像生成和分割等计算机视觉相关的多项任务，显示出其灵活性和强大能力。

ViT示例代码

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets, transforms
from torch.utils.data import DataLoader


# 定义 Patch Embedding 类，将输入图像划分为小块并转换为嵌入向量
class PatchEmbedding(nn.Module):
    def __init__(self, img_size=224, patch_size=16, in_channels=3, embed_dim=768):
        super(PatchEmbedding, self).__init__()
        self.patch_size = patch_size  # 每个小块的大小
        self.img_size = img_size  # 输入图像的大小
        self.embed_dim = embed_dim  # 嵌入的维度
        self.num_patches = (img_size // patch_size) ** 2  # 计算小块的数量

        # 使用卷积层将小块转换为嵌入向量
        self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size)

    def forward(self, x):
        # 通过卷积层提取小块特征
        x = self.proj(x)
        x = x.flatten(2)  # 将空间维度展平为二维数组
        x = x.transpose(1, 2)  # 转置，使得每个小块成为序列中的元素 (N, num_patches, embed_dim)
        return x


# 定义基本的 MLP Block，用于处理 Transformer 的前馈网络部分
class MLPBlock(nn.Module):
    def __init__(self, in_features, hidden_features, drop=0.1):
        super(MLPBlock, self).__init__()
        self.fc1 = nn.Linear(in_features, hidden_features)  # 第一层全连接
        self.dropout = nn.Dropout(drop)  # Dropout 层，防止过拟合
        self.fc2 = nn.Linear(hidden_features, in_features)  # 第二层全连接，将特征映射回输入维度

    def forward(self, x):
        # 前向传播：全连接 -> 激活 -> Dropout -> 全连接
        return self.fc2(self.dropout(F.relu(self.fc1(x))))


# 定义 Transformer 编码器的基本模块，包括自注意力和前馈网络
class EncoderBlock(nn.Module):
    def __init__(self, embed_dim, num_heads, drop=0.1):
        super(EncoderBlock, self).__init__()
        # 初始化多头自注意力机制
        self.attention = nn.MultiheadAttention(embed_dim, num_heads)
        self.ln1 = nn.LayerNorm(embed_dim)  # 第一层归一化
        self.ln2 = nn.LayerNorm(embed_dim)  # 第二层归一化
        self.drop_path1 = nn.Dropout(drop)  # Dropout 防止过拟合
        self.drop_path2 = nn.Dropout(drop)  # Dropout 防止过拟合
        self.mlp = MLPBlock(embed_dim, 4 * embed_dim, drop)  # 中间隐藏层维度为 4 倍于嵌入维度

    def forward(self, x):
        # 自注意力计算
        attn_output, _ = self.attention(x, x, x)  # 根据当前输入计算注意力得分
        # 残差连接 + 层归一化
        x = self.ln1(x + self.drop_path1(attn_output))
        # MLP 处理 + 残差连接 + 层归一化
        x = self.ln2(x + self.drop_path2(self.mlp(x)))
        return x


# 定义 Vision Transformer 模型
class VisionTransformer(nn.Module):
    def __init__(self, img_size=64, patch_size=16, num_classes=10, embed_dim=256, num_heads=4, num_layers=6):
        super(VisionTransformer, self).__init__()

        # 进行 Patch Embedding
        self.patch_embedding = PatchEmbedding(img_size, patch_size, in_channels=3, embed_dim=embed_dim)
        # 定义分类标记为可学习参数，形状为 (1, 1, embed_dim)
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        # 位置编码，形状为 (1, num_patches + 1, embed_dim)
        self.position_embeddings = nn.Parameter(torch.zeros(1, self.patch_embedding.num_patches + 1, embed_dim))
        # 堆叠多个 Transformer 编码器
        self.encoder_blocks = nn.ModuleList([EncoderBlock(embed_dim, num_heads) for _ in range(num_layers)])
        self.ln = nn.LayerNorm(embed_dim)  # 最后的层归一化
        self.mlp_head = nn.Linear(embed_dim, num_classes)  # 分类头

    def forward(self, x):
        batch_size = x.size(0)  # 获取当前 batch 的大小
        x = self.patch_embedding(x)  # 进行 Patch Embedding
        cls_tokens = self.cls_token.expand(batch_size, -1, -1)  # 扩展类标记的尺寸
        # 将类标记与嵌入特征拼接，并添加位置编码
        x = torch.cat((cls_tokens, x), dim=1) + self.position_embeddings

        # 通过所有的编码器块
        for block in self.encoder_blocks:
            x = block(x)

        x = self.ln(x)  # 最后的层归一化
        # 取类标记对应的输出，并通过分类头
        x = self.mlp_head(x[:, 0])  # 只能用 cls_token 的输出进行分类
        return x


# 定义训练过程
def train(model, dataloader, criterion, optimizer, device):
    model.train()  # 将模型设为训练模式
    total_loss = 0.0  # 用于累加损失
    num_batches = len(dataloader)  # 总批次数

    for images, labels in dataloader:
        images, labels = images.to(device), labels.to(device)  # 将数据转移到设备上

        optimizer.zero_grad()  # 清空梯度
        outputs = model(images)  # 前向传播
        loss = criterion(outputs, labels)  # 计算损失
        loss.backward()  # 反向传播
        optimizer.step()  # 更新参数

        total_loss += loss.item()  # 累加损失

    avg_loss = total_loss / num_batches  # 计算平均损失
    print(f'Average Loss: {avg_loss:.4f}')  # 打印平均损失


# 设置超参数 这里为了减少训练难度，降低了很多超参数的值
img_size = 64  # 输入图像的大小
batch_size = 16  # 每个批次的样本数量
learning_rate = 1e-4  # 学习率
num_epochs = 10  # 训练轮数
embed_dim = 256  # 嵌入维度
num_layers = 6  # Transformer的层数 
num_heads = 4  # 注意力头数

# 选择设备，优先使用 GPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 数据预处理和加载
transform = transforms.Compose([
    transforms.Resize((img_size, img_size)),  # 调整图像大小
    transforms.ToTensor(),  # 转换为张量
])

# 加载 CIFAR-10 数据集
train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)  # 数据加载器

# 初始化模型、损失函数和优化器
model = VisionTransformer(num_classes=10).to(device)  # 初始化 ViT 模型
criterion = nn.CrossEntropyLoss()  # 交叉熵损失
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)  # Adam 优化器

# 训练模型
for epoch in range(num_epochs):
    print(f'Starting Epoch {epoch + 1}/{num_epochs}')  # 打印当前 epoch
    train(model, train_loader, criterion, optimizer, device)  # 进行训练
    print(f'Epoch {epoch + 1}/{num_epochs} completed.')  # 打印完成状态

小结

在Vision Transformer (ViT)模型中，自注意力机制被引入并放置在卷积神经网络（CNN）之前，其具体功能和作用可以总结如下：
自注意力机制的作用功能：
全局上下文建模：自注意力机制能够捕捉图像中不同位置之间的关系，无论它们在空间上有多远。这使得网络能够全局考虑信息，从而提高对复杂模式和关系的建模能力。
特征重标定：自注意力机制通过计算各个特征之间的相似性，从而为特征赋予权重，有助于突出重要特征并抑制不重要的特征。
卷积网络：
局部特征捕捉：卷积网络通过局部卷积操作高效提取图像的局部特征，如边缘、纹理等。这些特征非常适合处理图像的空间局部性。
平移不变性：卷积操作具有平移不变性，即卷积层可以对输入图像的位移进行鲁棒处理，识别不同位置的相同物体。

除ViT之外，还有很多将自注意力机制放置在卷积网络之后的例子，实际上，无论自注意力机制和卷积网络的位置如何改变，自注意力机制一般都是起特征重标定，并通过全局信息增强特征的作用，而卷积神经网络因为自身特性，始终是起局部特征捕捉的作用。通过改变模型整体的局部结构可能会起到改善特定结果的作用，但是从这些模型的实际表现来看，进一步增强性能还需搭配如残差网络等其他结构才能同步增强模型的整体推理性能。

其他模型简介与示例

1. Non-Local Neural Networks (NLN)

基本概念: 非局部神经网络是一种在深度学习中引入非局部操作的框架，以便捕捉长距离依赖关系。传统的卷积神经网络主要关注局部邻域，而非局部神经网络可以考虑整个特征图上任意位置之间的关系。

工作原理: 各个位置的特征通过计算它们之间的相似性（通常使用点积或窗口计算）来生成逐位置的注意力权重。然后，这些权重被用来加权求和整个特征图，从而增强重要信息。

应用: NLN模型在目标检测和视频分析等任务中表现出色，因为它们能够更好地理解对象之间的关系和上下文。

2. Attention-Augmented Convolutional Networks

基本概念: 这类网络在传统卷积神经网络架构中引入了自注意力机制。这一设计结合了卷积的局部特征提取能力与自注意力的全局上下文信息建模能力。

工作原理: 在经过几层卷积层后，注意力模块会操作生成的特征图，以增强特征的表达。注意力机制主要在通道或空间维度上进行，从而使得模型能有效关注重要特征。

应用: 这种组合在很多视觉任务中表现良好，包括图像分类和目标检测，使得模型能够利用全局信息更好地理解图像内容。

3. Squeeze-and-Excitation Networks (SE)

基本概念: SE网络是通过引入一种特征重标定机制来增强卷积神经网络的性能的模型。它在传统卷积操作之后增加了注意力模块。

工作原理: SE模块通过“挤压”操作（全局平均池化）生成通道描述符，并通过“激励”操作（全连接层）自适应地为不同通道分配权重。这些权重用于重标定卷积特征，从而使得有用的特征得到增强，而不重要的特征则被压制。

应用: SE模块可以嵌入到多种卷积架构中，广泛应用于图像分类和图像生成任务等。

4. Combining CNN with Transformer Layers

基本概念: 一些现代神经网络架构（如EfficientNet和ResNet）结合了卷积神经网络与Transformer层。这种方法旨在利用卷积层的局部特征提取能力和Transformer层的全局上下文建模能力。

工作原理: 这些网络通常会在经过一系列卷积层提取特征后，引入Transformer层来处理提取出的特征图，增强网络对长距离依赖关系的捕捉能力。

应用: 这种混合结构在图像分类、目标检测和语义分割等任务中得到了很好的应用和验证，提高了模型的准确性与鲁棒性。

示例代码

自注意力模块放在卷积网络之后的示例

import torch
import torch.nn as nn
import torchvision.transforms as transforms
from torchvision import datasets
from torch.utils.data import DataLoader
import torch.optim as optim


# 定义多头自注意力机制
class MultiHeadAttention(nn.Module):
    def __init__(self, embed_dim, num_heads):
        super(MultiHeadAttention, self).__init__()
        self.num_heads = num_heads  # 注意力头的数量
        self.embed_dim = embed_dim  # 嵌入维度

        # 定义线性变换，用于获取值、键和查询
        self.values = nn.Linear(embed_dim, embed_dim, bias=False)  # 值线性变换
        self.keys = nn.Linear(embed_dim, embed_dim, bias=False)    # 键线性变换
        self.queries = nn.Linear(embed_dim, embed_dim, bias=False) # 查询线性变换
        self.fc_out = nn.Linear(embed_dim, embed_dim)  # 最终输出的线性变换

    def forward(self, x):
        N, sequence_length, embed_dim = x.shape  # 获取输入的批量大小、序列长度和嵌入维度
        values = self.values(x)  # 获取值
        keys = self.keys(x)      # 获取键
        queries = self.queries(x)  # 获取查询

        # 将每个头的值、键和查询分开
        values = values.view(N, sequence_length, self.num_heads, embed_dim // self.num_heads)
        keys = keys.view(N, sequence_length, self.num_heads, embed_dim // self.num_heads)
        queries = queries.view(N, sequence_length, self.num_heads, embed_dim // self.num_heads)

        # 重新排列以便计算
        values = values.permute(0, 2, 1, 3)  # (N, num_heads, sequence_length, embed_dim/num_heads)
        keys = keys.permute(0, 2, 1, 3)      # (N, num_heads, sequence_length, embed_dim/num_heads)
        queries = queries.permute(0, 2, 1, 3)  # (N, num_heads, sequence_length, embed_dim/num_heads)

        # 进行缩放点积自注意力计算
        energy = torch.einsum("nhqd,nhkd->nhqk", queries, keys)  # 计算量
        attention = torch.softmax(energy / (embed_dim ** (1 / 2)), dim=3)  # 计算注意力权重

        # 应用注意力权重
        out = torch.einsum("nhql,nhld->nhqd", attention, values)  # 计算输出
        out = out.permute(0, 2, 1, 3).contiguous().view(N, sequence_length, embed_dim)  # 变换输出维度
        return self.fc_out(out)  # 通过线性层得到最终输出


# 定义结合 CNN 和自注意力机制的网络
class CNNWithAttention(nn.Module):
    def __init__(self, num_classes):
        super(CNNWithAttention, self).__init__()
        # 定义卷积层
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)  # 第一卷积层，输入通道为3，输出通道为16
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)  # 第二卷积层，输入通道为16，输出通道为32
        self.pool = nn.MaxPool2d(2, 2)  # 最大池化层，2x2 的池化窗口
        self.fc1 = nn.Linear(32 * 8 * 8, 128)  # 全连接层，将特征映射到128维
        self.fc2 = nn.Linear(128, num_classes)  # 最后一个全连接层，将128维映射到类数
        self.attention = MultiHeadAttention(embed_dim=128, num_heads=4)  # 初始化多头自注意力机制

    def forward(self, x):
        # 前向传播
        x = self.pool(nn.ReLU()(self.conv1(x)))  # 卷积 -> ReLU 激活 -> 最大池化
        x = self.pool(nn.ReLU()(self.conv2(x)))  # 卷积 -> ReLU 激活 -> 最大池化
        x = x.view(x.size(0), -1)  # 展平操作，将多维输入展平为一维
        x = nn.ReLU()(self.fc1(x))  # 通过全连接层并应用 ReLU 激活

        # 自注意力机制
        x = x.unsqueeze(1)  # 增加时间维度，转换为 (N, 1, 128)
        x = self.attention(x)  # 应用多头自注意力机制
        x = x.squeeze(1)  # 移除时间维度，转换回 (N, 128)

        x = self.fc2(x)  # 最后通过分类层
        return x


# 数据预处理和加载
transform = transforms.Compose([
    transforms.Resize((32, 32)),  # 将图像大小调整为 32x32
    transforms.ToTensor(),  # 将图像转换为 Tensor 格式
])

# 下载并加载 CIFAR-10 数据集
train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)  # 下载训练数据
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)  # 数据加载器，批量大小为64，随机打乱

# 模型、损失和优化器
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  # 选择设备
model = CNNWithAttention(num_classes=10).to(device)  # 初始化模型并转移到指定设备
criterion = nn.CrossEntropyLoss()  # 选择交叉熵损失函数
optimizer = optim.Adam(model.parameters(), lr=0.001)  # Adam 优化器，学习率为0.001

# 训练模型
num_epochs = 20  # 设置训练轮数
for epoch in range(num_epochs):
    model.train()  # 设置模型为训练模式
    running_loss = 0.0  # 初始化当前轮的损失
    for images, labels in train_loader:  # 遍历训练加载器
        images, labels = images.to(device), labels.to(device)  # 将数据转移到设备

        optimizer.zero_grad()  # 清零梯度
        outputs = model(images)  # 前向传播获取模型输出
        loss = criterion(outputs, labels)  # 计算损失
        loss.backward()  # 反向传播计算梯度
        optimizer.step()  # 更新模型参数

        running_loss += loss.item()  # 累加当前批次的损失

    # 打印当前轮的平均损失
    print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {running_loss / len(train_loader):.4f}')

print("训练完成")  # 完成所有训练轮

训练结果

Epoch [1/20], Loss: 1.5621
Epoch [2/20], Loss: 1.1824
Epoch [3/20], Loss: 1.0215
Epoch [4/20], Loss: 0.9183
Epoch [5/20], Loss: 0.8360
Epoch [6/20], Loss: 0.7774
Epoch [7/20], Loss: 0.7236
Epoch [8/20], Loss: 0.6708
Epoch [9/20], Loss: 0.6310
Epoch [10/20], Loss: 0.5953
Epoch [11/20], Loss: 0.5564
Epoch [12/20], Loss: 0.5184
Epoch [13/20], Loss: 0.4822
Epoch [14/20], Loss: 0.4586
Epoch [15/20], Loss: 0.4209
Epoch [16/20], Loss: 0.4052
Epoch [17/20], Loss: 0.3763
Epoch [18/20], Loss: 0.3522
Epoch [19/20], Loss: 0.3377
Epoch [20/20], Loss: 0.3152
训练完成