补充:tensor之间进行矩阵相乘的方法总结

利用@进行简单的矩阵乘

@符号在tensor中就表示矩阵相乘,@符号的矩阵相乘性质在numpy中依然适用。

  • 首先矩阵相乘的双方必须满足可以矩阵相乘的条件
  • @只会关注两个矩阵最里面的两个维度是否符合条件,外面的维度都只表示矩阵运算的次数,甚至两个矩阵只要满足广播的条件和里面两个维度可以进行矩阵乘,二者的维度都可以不一样。

torch.mul

定要注意这个函数是陷阱!其与*的作用是完全一样的,其不管相乘的双方维度如何,执行的都是对位相乘的操作, *与torch.mul均不能实现矩阵相乘的规则。

torch.mm

torch.mm是阉割版的@,其只能对二维的tensor进行矩阵相乘,高了的维度其不会进行广播 ↓

a=torch.ones((2,3))
b=torch.ones((3,4))
print(torch.mm(a,b))
'''
tensor([[3., 3., 3., 3.],
        [3., 3., 3., 3.]])
'''

torch.matmul

其作用与@完全相同 ↓

a=torch.ones((1,2,2,3))
b=torch.ones((2,1,3,4))
print(torch.matmul(a,b))
'''
tensor([[[[3., 3., 3., 3.],
          [3., 3., 3., 3.]],

         [[3., 3., 3., 3.],
          [3., 3., 3., 3.]]],


        [[[3., 3., 3., 3.],
          [3., 3., 3., 3.]],

         [[3., 3., 3., 3.],
          [3., 3., 3., 3.]]]])
          '''

总结:就认准@就可以了

from dataclasses import dataclass from typing import Optional, Tuple Tensor = torch.Tensor @dataclass class AttentionOutput: hidden_states: Tensor attn_scores: Tensor class MultiHeadAttention(nn.Module): def __init__(self, config): super().__init__() # hyper params self.hidden_size = config["d_model"] # 隐藏层大小 self.num_heads = config["num_heads"] # 多头注意力的头数 assert ( self.hidden_size % self.num_heads == 0 ), "Hidden size must be divisible by num_heads but got {} and {}".format( self.hidden_size, self.num_heads ) self.head_dim = self.hidden_size // self.num_heads # 每个头的维度 # layers self.Wq = nn.Linear(self.hidden_size, self.hidden_size, bias=False) #第二个self.hidden_size可以*系数 self.Wk = nn.Linear(self.hidden_size, self.hidden_size, bias=False) self.Wv = nn.Linear(self.hidden_size, self.hidden_size, bias=False) self.Wo = nn.Linear(self.hidden_size, self.hidden_size, bias=False) # 输出层 def _split_heads(self, x: Tensor) -> Tensor: bs, seq_len, _ = x.shape #假设输入的维度是[batch_size, seq_len, hidden_size],hidden_size是512 x = x.view(bs, seq_len, self.num_heads, self.head_dim) #num_heads是8,head_dim是64 return x.permute(0, 2, 1, 3) #变换维度,[batch_size, num_heads, seq_len, head_dim] def _merge_heads(self, x: Tensor) -> Tensor:#将多头注意力的输出合并为一个张量 bs, _, seq_len, _ = x.shape #假设输入的维度是[batch_size, num_heads, seq_len, head_dim] return x.permute(0, 2, 1, 3).reshape(bs, seq_len, self.hidden_size) # 变换维度,变为[batch_size, seq_len, hidden_size] def forward(self, querys, keys, values, attn_mask=None) -> AttentionOutput: # split heads querys = self._split_heads(self.Wq(querys)) #(batch_size, seq_len,hidden_dim)-->[batch_size, num_heads, seq_len, head_dim] keys = self._split_heads(self.Wk(keys))#[batch_size, num_heads, seq_len, head_dim] values = self._split_heads(self.Wv(values))#[batch_size, num_heads, seq_len, head_dim] # calculate attention scores qk_logits = torch.matmul(querys, keys.mT) # 计算注意力分数,matmul是矩阵乘法,mT是矩阵转置,qk_logits是[batch_size, num_heads, seq_len, seq_len] # print(querys.shape[-2], keys.shape[-2]) #3 4 if attn_mask is not None: attn_mask = attn_mask[:, :, : querys.shape[-2], : keys.shape[-2]] qk_logits += attn_mask * -1e9 # 给需要mask的地方设置一个负无穷 attn_scores = F.softmax(qk_logits / (self.head_dim**0.5), dim=-1) # 计算注意力分数 # embeds的尺寸是[batch_size, num_heads, seq_len, head_dim] embeds = torch.matmul(attn_scores, values) # softmax后的结果与value相乘,得到新的表示 embeds = self.Wo(self._merge_heads(embeds)) # 输出层 [batch_size, seq_len, hidden_size] return AttentionOutput(hidden_states=embeds, attn_scores=attn_scores) mha = MultiHeadAttention({"num_heads": 2, "d_model": 2}) query = torch.randn(2, 3, 2) # [batch_size, seq_len, hidden_size] query /= query.norm(dim=-1, keepdim=True) # 归一化 key_value = torch.randn(2, 4, 2) print(f'key_value.shape {key_value.shape}') outputs = mha(query, key_value, key_value) #最终输出shape和query的shape一样 print(outputs.hidden_states.shape) print(outputs.attn_scores.shape)这段代码有什么用
03-27
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值