PVT的spatial reduction attention(SRA)

就是用了一个卷积降了一下k,v 的size

可以理解为将R个点聚合成一个，然后attention的时候Q和聚合成的点的K和V算

import torch
from torch import nnclass SpatialReductionAttention(nn.Module):def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0., sr_ratio=1):super().__init__()assert dim % num_heads == 0, f"dim {dim} should be divided by num_heads {num_heads}."self.dim = dimself.num_heads = num_headshead_dim = dim // num_headsself.scale = qk_scale or head_dim ** -0.5self.q = nn.Linear(dim, dim, bias=qkv_bias)self.kv = nn.Linear(dim, dim * 2, bias=qkv_bias)self.attn_drop = nn.Dropout(attn_drop)self.proj = nn.Linear(dim, dim)self.dropout = nn.Dropout(proj_drop)self.sr_ratio = sr_ratio# 实现上这里等价于一个卷积层if sr_ratio > 1:self.sr = nn.Conv2d(dim, dim, kernel_size=sr_ratio, stride=sr_ratio)self.norm = nn.LayerNorm(dim)def forward(self, x, H, W):B, N, D = x.shape  #N=h*wq = self.q(x).reshape(B, N, self.num_heads, D // self.num_heads).permute(0, 2, 1, 3)if self.sr_ratio > 1:x_ = x.permute(0, 2, 1).reshape(B, D, H, W)x_ = self.sr(x_).reshape(B, D, -1).permute(0, 2, 1) # 这里x_.shape = (B, N/R^2, D)x_ = self.norm(x_)kv = self.kv(x_).reshape(B, -1, 2, self.num_heads, D // self.num_heads).permute(2, 0, 3, 1, 4)else:kv = self.kv(x).reshape(B, -1, 2, self.num_heads, D // self.num_heads).permute(2, 0, 3, 1, 4)k, v = kv[0], kv[1]attn = (q @ k.transpose(-2, -1)) * self.scaleattn = attn.softmax(dim=-1)attn = self.attn_drop(attn)x = (attn @ v).transpose(1, 2).reshape(B, N, D)x = self.proj(x)x = self.dropout(x)return xx = torch.rand(4, 224*128, 256)
attn = SpatialReductionAttention(dim=256, sr_ratio = 2)
output = attn(x, H=224, W=128)

PVT的spatial reduction attention(SRA)相关推荐

【论文笔记】SPAN: Spatial Pyramid Attention Network for Image Manipulation Localization
SPAN: Spatial Pyramid Attention Network for Image Manipulation Localization 发布于ECCV2020 原文链接:https:/ ...
Vision Transformer在CV任务中的速度如何保证？
本文作者丨盘子正@知乎编辑丨极市平台来源丨https://zhuanlan.zhihu.com/p/569482746 我(盘子正@知乎)的PhD课题是Vision Transformer的 ...
Transformer合集1
最近Transformer文章太多了索性一起发了得~~ 以后关于这个的都不单发了如何提高ViT的效率?可以是让模型更容易训练,减少训练时间,也可以减少模型部署在硬件上的功耗等等.本文主要讲inf ...
SepViT：可分离视觉Transformer
点击下方卡片,关注"CVer"公众号 AI/CV重磅干货,第一时间送达转载自:集智书童 SepViT: Separable Vision Transformer 论文:https ...
论文精读:PVT v2: Improved Baselines with Pyramid Vision Transformer
论文地址:https://arxiv.org/abs/2106.13797 源码地址:https://github.com/whai362/PVT Abstract 在这项工作中,作者改进了PVT v ...
论文：Pyramid Vision Transformer
Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions 金字塔视觉Tran ...
PVTV2--Pyramid Vision TransformerV2学习笔记
PVTV2–Pyramid Vision TransformerV2学习笔记 PVTv2: Improved Baselines with Pyramid Vision Transformer Abs ...
【深度学习】（ICCV-2021）PVT-金字塔 Vision Transformer及PVT_V2
目录 0. 详情 1. 简述 2.主要工作 2.1 ViT遗留的问题 2.2 引入金字塔结构 3.PVT的设计方案 3.1 Patch embedding 代码 3.2position embeddi ...
transformer与视觉
目录综述优秀网文基本transformer 视觉transformer原理具体的transformer 一般方法 ViT :一张图等于 16x16 个字,计算机视觉也用上 Transforme ...

PVT的spatial reduction attention(SRA)

PVT的spatial reduction attention(SRA)相关推荐

最新文章

热门文章