torch.bmm()函数详解

原创已于 2023-07-13 17:07:25 修改

· 479 阅读

1 ·

版权

文章标签：

#python #深度学习 #机器学习

于 2023-07-13 17:06:03 首次发布

计算机视觉专栏收录该内容

14 篇文章

订阅专栏

该代码定义了一个名为PAM_Module的神经网络模块，用于处理输入特征映射。它包含查询、键和值的卷积层，通过矩阵乘法计算注意力权重，并使用softmax函数归一化。最后，将注意力加权值与输入特征相加得到输出。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

去掉batch_size这一维度，按照矩阵乘法运算即可

自行体会

class PAM_Module(nn.Module):
    """ Position attention module"""
    def __init__(self, in_dim) -> None:
        super().__init__()
        self.channel_in = in_dim
        
        self.query_conv = Conv2d(in_dim, in_dim // 8, kernel_size=1)
        self.key_conv = Conv2d(in_dim, in_dim//8, kernel_size=1)
        self.value_conv = Conv2d(in_dim, in_dim, kernel_size=1)
        
        self.gamma = Parameter(torch.zeros(1))
        
        self.softmax = Softmax(dim = -1)
        
    def forward(self, x):
        """
            inputs :
                x : input feature maps( B X C X H X W)
            returns :
                out : attention value + input feature
                attention: B X (HxW) X (HxW)
        """
        ## B C H W
        m_batchsize, C, height, width = x.size()
        ## B C H W --> B C HW --> B HW C
        proj_query = self.query_conv(x).view(m_batchsize, -1, height*width).permute(0, 2, 1)
        ## B C HW
        proj_key = self.key_conv(x).view(m_batchsize, -1, height*width)
        ## B HW HW
        energy = torch.bmm(proj_query, proj_key)
        attention = self.softmax(energy)
        ## B C HW 
        proj_value = self.value_conv(x).view(m_batchsize, -1, height*width)
        
        ## bmm(B C HW, B HW HW) --> B C HW
        out = torch.bmm(proj_value, attention.permute(0, 2, 1))
        ## B C H W
        out = out.view(m_batchsize, C, height, width)
        ## B C H W
        out = self.gamma*out + x
        
        return out