注意力机制和自注意力机制模块的相关代码实现/缝合模块/即插即用模块

复制代码
注意力机制

import torch
import torch.nn as nn

class Attention(nn.Module):
    def __init__(self, hidden_dim):
        super(Attention, self).__init__()
        self.attention = nn.Linear(hidden_dim, 1, bias=False)

    def forward(self, encoder_outputs):
        # encoder_outputs shape: (batch_size, sequence_length, hidden_dim)
        attn_weights = self.attention(encoder_outputs)  # (batch_size, sequence_length, 1)
        attn_weights = torch.softmax(attn_weights, dim=1)  # (batch_size, sequence_length, 1)
        context = torch.sum(attn_weights * encoder_outputs, dim=1)  # (batch_size, hidden_dim)
        return context, attn_weights

# 示例用法
batch_size = 2
sequence_length = 5
hidden_dim = 10

encoder_outputs = torch.randn(batch_size, sequence_length, hidden_dim)
attention_layer = Attention(hidden_dim)
context, attn_weights = attention_layer(encoder_outputs)

print("Context:", context)
print("Attention Weights:", attn_weights)

自注意力机制

import torch

import torch.nn as nn

class MultiHeadSelfAttention(nn.Module):

def init(self, embed_dim, num_heads):

super(MultiHeadSelfAttention, self).init()

self.embed_dim = embed_dim

self.num_heads = num_heads

self.head_dim = embed_dim // num_heads

assert self.head_dim * num_heads == embed_dim, "embed_dim must be divisible by num_heads"

self.qkv_proj = nn.Linear(embed_dim, embed_dim * 3)

self.o_proj = nn.Linear(embed_dim, embed_dim)

self.softmax = nn.Softmax(dim=-1)

def forward(self, x):

batch_size, seq_length, embed_dim = x.size()

qkv = self.qkv_proj(x) # (batch_size, seq_length, embed_dim * 3)

qkv = qkv.reshape(batch_size, seq_length, self.num_heads, 3 * self.head_dim)

qkv = qkv.permute(0, 2, 1, 3) # (batch_size, num_heads, seq_length, 3 * head_dim)

q, k, v = qkv.chunk(3, dim=-1) # Each has shape (batch_size, num_heads, seq_length, head_dim)

attn_weights = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5) # Scaled dot-product

attn_weights = self.softmax(attn_weights) # (batch_size, num_heads, seq_length, seq_length)

attn_output = torch.matmul(attn_weights, v) # (batch_size, num_heads, seq_length, head_dim)

attn_output = attn_output.permute(0, 2, 1, 3).contiguous()

attn_output = attn_output.reshape(batch_size, seq_length, embed_dim)

output = self.o_proj(attn_output)

return output, attn_weights

示例用法

batch_size = 2

seq_length = 5

embed_dim = 16

num_heads = 4

x = torch.randn(batch_size, seq_length, embed_dim)

self_attention_layer = MultiHeadSelfAttention(embed_dim, num_heads)

output, attn_weights = self_attention_layer(x)

print("Output:", output)

print("Attention Weights:", attn_weights)

相关推荐
秀儿还能再秀6 小时前
神经网络(系统性学习三):多层感知机(MLP)
神经网络·学习笔记·mlp·多层感知机
老艾的AI世界8 小时前
AI翻唱神器,一键用你喜欢的歌手翻唱他人的曲目(附下载链接)
人工智能·深度学习·神经网络·机器学习·ai·ai翻唱·ai唱歌·ai歌曲
sp_fyf_202411 小时前
【大语言模型】ACL2024论文-19 SportsMetrics: 融合文本和数值数据以理解大型语言模型中的信息融合
人工智能·深度学习·神经网络·机器学习·语言模型·自然语言处理
CoderIsArt11 小时前
基于 BP 神经网络整定的 PID 控制
人工智能·深度学习·神经网络
z千鑫12 小时前
【人工智能】PyTorch、TensorFlow 和 Keras 全面解析与对比:深度学习框架的终极指南
人工智能·pytorch·深度学习·aigc·tensorflow·keras·codemoss
EterNity_TiMe_12 小时前
【论文复现】神经网络的公式推导与代码实现
人工智能·python·深度学习·神经网络·数据分析·特征分析
思通数科多模态大模型13 小时前
10大核心应用场景,解锁AI检测系统的智能安全之道
人工智能·深度学习·安全·目标检测·计算机视觉·自然语言处理·数据挖掘
数据岛13 小时前
数据集论文:面向深度学习的土地利用场景分类与变化检测
人工智能·深度学习
学不会lostfound13 小时前
三、计算机视觉_05MTCNN人脸检测
pytorch·深度学习·计算机视觉·mtcnn·p-net·r-net·o-net
红色的山茶花13 小时前
YOLOv8-ultralytics-8.2.103部分代码阅读笔记-block.py
笔记·深度学习·yolo