每日Attention学习15——Cross-Model Grafting Module

模块出处

[CVPR 22] [link] [code] Pyramid Grafting Network for One-Stage High Resolution Saliency Detection


模块名称

Cross-Model Grafting Module (CMGM)


模块作用

Transformer与CNN之间的特征融合


模块结构

模块思想

Transformer在全局特征上更优,CNN在局部特征上更优,对这两者进行进行融合的最简单做法是直接相加或相乘。但是,相加或相乘本质上属于"局部"操作,如果某片区域两个特征的不确定性都较高,则会带来许多噪声。为此,本文提出了CMGM模块,通过交叉注意力的形式引入更为广泛的信息来增强融合效果。


模块代码
python 复制代码
import torch.nn.functional as F
import torch.nn as nn
import torch


class CMGM(nn.Module):
    def __init__(self, dim, num_heads=8, qkv_bias=True, qk_scale=None):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = qk_scale or head_dim ** -0.5
        self.k = nn.Linear(dim, dim , bias=qkv_bias)
        self.qv = nn.Linear(dim, dim * 2, bias=qkv_bias)
        self.proj = nn.Linear(dim, dim)
        self.act = nn.ReLU(inplace=True)
        self.conv = nn.Conv2d(8,8,kernel_size=3, stride=1, padding=1)
        self.lnx = nn.LayerNorm(64)
        self.lny = nn.LayerNorm(64)
        self.bn = nn.BatchNorm2d(8)
        self.conv2 = nn.Sequential(
            nn.Conv2d(64,64,kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64,64,kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True)
        )
        
    def forward(self, x, y):
        batch_size = x.shape[0]
        chanel     = x.shape[1]
        sc = x
        x = x.view(batch_size, chanel, -1).permute(0, 2, 1)
        sc1 = x
        x = self.lnx(x)
        y = y.view(batch_size, chanel, -1).permute(0, 2, 1)
        y = self.lny(y)
        
        B, N, C = x.shape
        y_k = self.k(y).reshape(B, N, 1, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        x_qv= self.qv(x).reshape(B,N,2,self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        x_q, x_v = x_qv[0], x_qv[1] 
        y_k = y_k[0]
        attn = (x_q @ y_k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)
        x = (attn @ x_v).transpose(1, 2).reshape(B, N, C)

        x = self.proj(x)
        x = (x+sc1)

        x = x.permute(0,2,1)
        x = x.view(batch_size,chanel,*sc.size()[2:])
        x = self.conv2(x)+x
        return x, self.act(self.bn(self.conv(attn+attn.transpose(-1,-2))))
    

if __name__ == '__main__':
    x = torch.randn([1, 64, 11, 11])
    y = torch.randn([1, 64, 11, 11])
    cmgm = CMGM(dim=64)
    out1, out2 = cmgm(x, y)
    print(out1.shape)  # out feature 1, 64, 11, 11
    print(out2.shape)  # cross attention matrix 1, 8, 121, 121

相关推荐
LuH11243 小时前
【论文阅读笔记】Scalable, Detailed and Mask-Free Universal Photometric Stereo
论文阅读·笔记
小嗷犬3 小时前
【论文笔记】Cross-lingual few-shot sign language recognition
论文阅读·人工智能·多模态·少样本·手语翻译
实验室里哈啤酒15 小时前
ResEmoteNet论文阅读与推理
论文阅读
LuH112417 小时前
【论文阅读笔记】Learning to sample
论文阅读·笔记·图形渲染·点云
开心星人21 小时前
【论文阅读】Unlearning Backdoor Attacks in Federated Learning
论文阅读
七夜星七夜月3 天前
时间序列预测论文阅读和相关代码库
论文阅读·python·深度学习
WenBoo-3 天前
HIPT论文阅读
论文阅读
chnyi6_ya3 天前
论文笔记:Buffer of Thoughts: Thought-Augmented Reasoning with Large Language Models
论文阅读·人工智能·语言模型
Jude_lennon4 天前
【论文笔记】结合:“integrate“ 和 “combine“等
论文阅读
LuH11244 天前
【论文阅读笔记】HunyuanVideo: A Systematic Framework For Large Video Generative Models
论文阅读·笔记