深度学习算法中的基于注意力机制的神经网络（Attention-based Neural Networks）

pythonCopy codeimport torch
import torch.nn as nn
import torch.nn.functional as F
class Attention(nn.Module):
    def __init__(self, hidden_size):
        super(Attention, self).__init__()
        self.hidden_size = hidden_size
        self.attn = nn.Linear(hidden_size, hidden_size)
    def forward(self, encoder_outputs, hidden):
        # encoder_outputs: [batch_size, seq_len, hidden_size]
        # hidden: [batch_size, hidden_size]
        seq_len = encoder_outputs.size(1)
        # 计算注意力权重
        attn_weights = F.softmax(
            self.attn(hidden).unsqueeze(1).bmm(encoder_outputs.transpose(1, 2)),
            dim=2
        )
        
        # 加权求和得到上下文向量
        context = attn_weights.bmm(encoder_outputs)
        
        return context, attn_weights
class AttentionBasedNetwork(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(AttentionBasedNetwork, self).__init__()
        self.hidden_size = hidden_size
        self.embedding = nn.Embedding(input_size, hidden_size)
        self.gru = nn.GRU(hidden_size, hidden_size)
        self.attention = Attention(hidden_size)
        self.out = nn.Linear(hidden_size, output_size)
    def forward(self, input_seq):
        embedded = self.embedding(input_seq)
        output, hidden = self.gru(embedded)
        # 使用注意力机制获取上下文向量
        context, attn_weights = self.attention(output, hidden.squeeze(0))
        # 将上下文向量和隐藏状态连接起来
        output = torch.cat((context, hidden.squeeze(0)), dim=1)
        output = self.out(output)
        # 使用softmax函数获取分类概率
        output = F.softmax(output, dim=1)
        return output, attn_weights
# 示例用法
input_size = 1000
hidden_size = 128
output_size = 10
seq_len = 20
batch_size = 32
# 创建模型实例
model = AttentionBasedNetwork(input_size, hidden_size, output_size)
# 构造输入数据
input_seq = torch.randint(input_size, (batch_size, seq_len))
# 前向传播
output, attn_weights = model(input_seq)
# 打印输出结果
print("输出概率：", output)
print("注意力权重：", attn_weights)

pythonCopy codeimport torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
# 定义注意力机制的模块
class Attention(nn.Module):
    def __init__(self, hidden_size):
        super(Attention, self).__init__()
        self.hidden_size = hidden_size
        self.attn = nn.Linear(hidden_size, hidden_size)
    def forward(self, encoder_outputs, hidden):
        seq_len = encoder_outputs.size(1)
        # 计算注意力权重
        attn_weights = torch.bmm(self.attn(hidden).unsqueeze(1), encoder_outputs.transpose(1, 2))
        attn_weights = torch.softmax(attn_weights, dim=2)
        # 加权求和得到上下文向量
        context = torch.bmm(attn_weights, encoder_outputs)
        return context.squeeze(1), attn_weights.squeeze(1)
# 定义基于注意力机制的图像分类网络
class AttentionBasedImageClassifier(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(AttentionBasedImageClassifier, self).__init__()
        self.hidden_size = hidden_size
        self.attention = Attention(hidden_size)
        self.fc = nn.Linear(hidden_size, output_size)
    def forward(self, input):
        # input的shape: [batch_size, channel, height, width]
        batch_size, channel, height, width = input.size()
        input = input.view(batch_size, channel, height * width)
        
        # 使用注意力机制获取上下文向量
        context, attn_weights = self.attention(input, input.mean(dim=2))
        
        # 全连接层进行分类
        output = self.fc(context)
        return output, attn_weights
# 数据预处理
transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5,), (0.5,))])
# 加载训练集和测试集
trainset = torchvision.datasets.MNIST(root="./data", train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=32, shuffle=True)
testset = torchvision.datasets.MNIST(root="./data", train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=32, shuffle=False)
# 创建模型实例
model = AttentionBasedImageClassifier(input_size=28*28, hidden_size=128, output_size=10)
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
# 训练模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
for epoch in range(5):
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data[0].to(device), data[1].to(device)
        optimizer.zero_grad()
        outputs, _ = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        if i % 200 == 199:
            print('[%d, %5d] loss: %.3f' %
                  (epoch + 1, i + 1, running_loss / 200))
            running_loss = 0.0
print("Finished training")
# 在测试集上评估模型
correct = 0
total = 0
with torch.no_grad():
    for data in testloader:
        images, labels = data[0].to(device), data[1].to(device)
        outputs, _ = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
print('Accuracy on the test images: %.2f %%' % (100 * correct / total))