1、初始版本
上图N=H*W
来自csnln中的code
class NonLocalAttention(nn.Module):
def __init__(self, channel=128, reduction=2, ksize=3, scale=3, stride=1, softmax_scale=10, average=True, conv=common.default_conv):
super(NonLocalAttention, self).__init__()
self.conv_match1 = common.BasicBlock(conv, channel, channel//reduction, 1, bn=False, act=nn.PReLU())
self.conv_match2 = common.BasicBlock(conv, channel, channel//reduction, 1, bn=False, act = nn.PReLU())
self.conv_assembly = common.BasicBlock(conv, channel, channel, 1,bn=False, act=nn.PReLU())
def forward(self, input):
x_embed_1 = self.conv_match1(input)
x_embed_2 = self.conv_match2(input)
x_assembly = self.conv_assembly(input)
N,C,H,W = x_embed_1.shape
x_embed_1 = x_embed_1.permute(0,2,3,1).view((N,H*W,C))
x_embed_2 = x_embed_2.view(N,C,H*W)
score = torch.matmul(x_embed_1, x_embed_2)
score = F.softmax(score, dim=2)
x_assembly = x_assembly.view(N,-1,H*W).permute(0,2,1)
x_final = torch.matmul(score, x_assembly)
return x_final.permute(0,2,1).view(N,-1,H,W)2、SAN中的分割成4份
①san中采用原始的non-local 论文提供的代码【CV中的Attention机制】Non-Local neural networks的理解与实现 - pprp - 博客园1. Non local Non Local是王小龙在CVPR2018年提出的一个自注意力模型。Non Local NN和Non Local Means非局部均值去燥滤波有点相似的感觉。普通的滤波都是
https://www.cnblogs.com/pprp/p/12153255.html
可以看出，第一个CSNLN是采用的embedded Gaussian来做的
class _NonLocalBlockND(nn.Module):
def __init__(self, in_channels, inter_channels=None, dimension=3, mode='embedded_gaussian',
sub_sample=True, bn_layer=True):
super(_NonLocalBlockND, self).__init__()
assert dimension in [1, 2, 3]
assert mode in ['embedded_gaussian', 'gaussian', 'dot_product', 'concatenation']
# print('Dimension: %d, mode: %s' % (dimension, mode))
self.mode = mode
self.dimension = dimension
self.sub_sample = sub_sample
self.in_channels = in_channels
self.inter_channels = inter_channels
if self.inter_channels is None:
self.inter_channels = in_channels // 2
if self.inter_channels == 0:
self.inter_channels = 1
if dimension == 3:
conv_nd = nn.Conv3d
max_pool = nn.MaxPool3d
bn = nn.BatchNorm3d
elif dimension == 2:
conv_nd = nn.Conv2d
max_pool = nn.MaxPool2d
sub_sample = nn.Upsample
bn = nn.BatchNorm2d
else:
conv_nd = nn.Conv1d
max_pool = nn.MaxPool1d
bn = nn.BatchNorm1d
self.g = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,
kernel_size=1, stride=1, padding=0)
if bn_layer:
self.W = nn.Sequential(
conv_nd(in_channels=self.inter_channels, out_channels=self.in_channels,
kernel_size=1, stride=1, padding=0),
bn(self.in_channels)
)
nn.init.constant_(self.W[1].weight, 0)
nn.init.constant_(self.W[1].bias, 0)
else:
self.W = conv_nd(in_channels=self.inter_channels, out_channels=self.in_channels,
kernel_size=1, stride=1, padding=0)
nn.init.constant_(self.W.weight, 0)
nn.init.constant_(self.W.bias, 0)
self.theta = None
self.phi = None
self.concat_project = None
# self.fc = nn.Linear(64,2304,bias=True)
# self.sub_bilinear = nn.Upsample(size=(48,48),mode='bilinear')
# self.sub_maxpool = nn.AdaptiveMaxPool2d(output_size=(48,48))
if mode in ['embedded_gaussian', 'dot_product', 'concatenation']:
self.theta = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,
kernel_size=1, stride=1, padding=0)
self.phi = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,
kernel_size=1, stride=1, padding=0)
if mode == 'embedded_gaussian':
self.operation_function = self._embedded_gaussian
elif mode == 'dot_product':
self.operation_function = self._dot_product
elif mode == 'concatenation':
self.operation_function = self._concatenation
self.concat_project = nn.Sequential(
nn.Conv2d(self.inter_channels * 2, 1, 1, 1, 0, bias=False),
nn.ReLU()
)
elif mode == 'gaussian':
self.operation_function = self._gaussian
if sub_sample:
self.g = nn.Sequential(self.g, max_pool(kernel_size=2))
if self.phi is None:
self.phi = max_pool(kernel_size=2)
else:
self.phi = nn.Sequential(self.phi, max_pool(kernel_size=2))
def forward(self, x):
'''
:param x: (b, c, t, h, w)
:return:
'''
output = self.operation_function(x)
return output
def _embedded_gaussian(self, x):
batch_size,C,H,W = x.shape
# x_sub = self.sub_bilinear(x) # bilinear downsample
# x_sub = self.sub_maxpool(x) # maxpool downsample
##
# g_x = x.view(batch_size, self.inter_channels, -1)
# g_x = g_x.permute(0, 2, 1)
#
# # theta=>(b, c, t, h, w)[->(b, 0.5c, t, h, w)]->(b, thw, 0.5c)
# # phi =>(b, c, t, h, w)[->(b, 0.5c, t, h, w)]->(b, 0.5c, thw)
# # f=>(b, thw, 0.5c)dot(b, 0.5c, twh) = (b, thw, thw)
# theta_x = x.view(batch_size, self.inter_channels, -1)
# theta_x = theta_x.permute(0, 2, 1)
# fc = self.fc(theta_x)
# # phi_x = self.phi(x).view(batch_size, self.inter_channels, -1)
# # f = torch.matmul(theta_x, phi_x)
# # return f
# # f_div_C = F.softmax(fc, dim=-1)
# return fc
##
# g=>(b, c, t, h, w)->(b, 0.5c, t, h, w)->(b, thw, 0.5c)
g_x = self.g(x).view(batch_size, self.inter_channels, -1)
g_x = g_x.permute(0, 2, 1)
# theta=>(b, c, t, h, w)[->(b, 0.5c, t, h, w)]->(b, thw, 0.5c)
# phi =>(b, c, t, h, w)[->(b, 0.5c, t, h, w)]->(b, 0.5c, thw)
# f=>(b, thw, 0.5c)dot(b, 0.5c, twh) = (b, thw, thw)
theta_x = self.theta(x).view(batch_size, self.inter_channels, -1)
theta_x = theta_x.permute(0, 2, 1)
phi_x = self.phi(x).view(batch_size, self.inter_channels, -1)
f = torch.matmul(theta_x, phi_x)
# return f
f_div_C = F.softmax(f, dim=-1)
# return f_div_C
# (b, thw, thw)dot(b, thw, 0.5c) = (b, thw, 0.5c)->(b, 0.5c, t, h, w)->(b, c, t, h, w)
y = torch.matmul(f_div_C, g_x)
y = y.permute(0, 2, 1).contiguous()
y = y.view(batch_size, self.inter_channels, *x.size()[2:])
W_y = self.W(y)
z = W_y + x
return z
def _gaussian(self, x):
batch_size = x.size(0)
g_x = self.g(x).view(batch_size, self.inter_channels, -1)
g_x = g_x.permute(0, 2, 1)
theta_x = x.view(batch_size, self.in_channels, -1)
theta_x = theta_x.permute(0, 2, 1)
if self.sub_sample:
phi_x = self.phi(x).view(batch_size, self.in_channels, -1)
else:
phi_x = x.view(batch_size, self.in_channels, -1)
f = torch.matmul(theta_x, phi_x)
f_div_C = F.softmax(f, dim=-1)
y = torch.matmul(f_div_C, g_x)
y = y.permute(0, 2, 1).contiguous()
y = y.view(batch_size, self.inter_channels, *x.size()[2:])
W_y = self.W(y)
z = W_y + x
return z
def _dot_product(self, x):
batch_size = x.size(0)
g_x = self.g(x).view(batch_size, self.inter_channels, -1)
g_x = g_x.permute(0, 2, 1)
theta_x = self.theta(x).view(batch_size, self.inter_channels, -1)
theta_x = theta_x.permute(0, 2, 1)
phi_x = self.phi(x).view(batch_size, self.inter_channels, -1)
f = torch.matmul(theta_x, phi_x)
N = f.size(-1)
f_div_C = f / N
y = torch.matmul(f_div_C, g_x)
y = y.permute(0, 2, 1).contiguous()
y = y.view(batch_size, self.inter_channels, *x.size()[2:])
W_y = self.W(y)
z = W_y + x
return z
def _concatenation(self, x):
batch_size = x.size(0)
g_x = self.g(x).view(batch_size, self.inter_channels, -1)
g_x = g_x.permute(0, 2, 1)
# (b, c, N, 1)
theta_x = self.theta(x).view(batch_size, self.inter_channels, -1, 1)
# (b, c, 1, N)
phi_x = self.phi(x).view(batch_size, self.inter_channels, 1, -1)
h = theta_x.size(2)
w = phi_x.size(3)
theta_x = theta_x.repeat(1, 1, 1, w)
phi_x = phi_x.repeat(1, 1, h, 1)
concat_feature = torch.cat([theta_x, phi_x], dim=1)
f = self.concat_project(concat_feature)
b, _, h, w = f.size()
f = f.view(b, h, w)
N = f.size(-1)
f_div_C = f / N
y = torch.matmul(f_div_C, g_x)
y = y.permute(0, 2, 1).contiguous()
y = y.view(batch_size, self.inter_channels, *x.size()[2:])
W_y = self.W(y)
z = W_y + x
return z②他分割成了四分：
feat_sub_lu = x[:, :, :H1, :W1]
feat_sub_ld = x[:, :, H1:, :W1]
feat_sub_ru = x[:, :, :H1, W1:]
feat_sub_rd = x[:, :, H1:, W1:]直接切片就是了
3、CSNLN中的cross scale版
class CrossScaleAttention(nn.Module):
def __init__(self, channel=128, reduction=2, ksize=3, scale=3, stride=1, softmax_scale=10, average=True, conv=common.default_conv):
super(CrossScaleAttention, self).__init__()
self.ksize = ksize
self.stride = stride
self.softmax_scale = softmax_scale
self.scale = scale
self.average = average
escape_NaN = torch.FloatTensor([1e-4])
self.register_buffer('escape_NaN', escape_NaN)
self.conv_match_1 = common.BasicBlock(conv, channel, channel//reduction, 1, bn=False, act=nn.PReLU())
self.conv_match_2 = common.BasicBlock(conv, channel, channel//reduction, 1, bn=False, act=nn.PReLU())
self.conv_assembly = common.BasicBlock(conv, channel, channel, 1, bn=False, act=nn.PReLU())
#self.register_buffer('fuse_weight', fuse_weight)
def forward(self, input):
#get embedding
embed_w = self.conv_assembly(input)
shape_input = list(embed_w.size()) # b*c*h*w
match_input = self.conv_match_1(input)
# b*c*h*w
input_groups = torch.split(match_input,1,dim=0)
# kernel size on input for matching
kernel = self.scale*self.ksize
# raw_w is extracted for reconstruction
raw_w = extract_image_patches(embed_w, ksizes=[kernel, kernel],
strides=[self.stride*self.scale,self.stride*self.scale],
rates=[1, 1],
padding='same') # [N, C*k*k, L]
# raw_shape: [N, C, k, k, L]
raw_w = raw_w.view(shape_input[0], shape_input[1], kernel, kernel, -1)
raw_w = raw_w.permute(0, 4, 1, 2, 3) # raw_shape: [N, L, C, k, k]
raw_w_groups = torch.split(raw_w, 1, dim=0)
# downscaling X to form Y for cross-scale matching
ref = F.interpolate(input, scale_factor=1./self.scale, mode='bilinear')
ref = self.conv_match_2(ref)
w = extract_image_patches(ref, ksizes=[self.ksize, self.ksize],
strides=[self.stride, self.stride],
rates=[1, 1],
padding='same')
shape_ref = ref.shape#(b,c,h/3,w/3)
# w shape: [N, C, k, k, L]
w = w.view(shape_ref[0], shape_ref[1], self.ksize, self.ksize, -1)
w = w.permute(0, 4, 1, 2, 3) # w shape: [N, L, C, k, k]
w_groups = torch.split(w, 1, dim=0)
y = []
scale = self.softmax_scale
# 1*1*k*k
#fuse_weight = self.fuse_weight
for xi, wi, raw_wi in zip(input_groups, w_groups, raw_w_groups):
# normalize
wi = wi[0] # [L, C, k, k]
max_wi = torch.max(torch.sqrt(reduce_sum(torch.pow(wi, 2),
axis=[1, 2, 3],
keepdim=True)),
self.escape_NaN)
wi_normed = wi/ max_wi
# Compute correlation map
xi = same_padding(xi, [self.ksize, self.ksize], [1, 1], [1, 1]) # xi: 1*c*H*W
yi = F.conv2d(xi, wi_normed, stride=1) # [1, L, H, W] L = shape_ref[2]*shape_ref[3]
yi = yi.view(1,shape_ref[2] * shape_ref[3], shape_input[2], shape_input[3]) # (B=1, C=32*32, H=32, W=32)
# rescale matching score
yi = F.softmax(yi*scale, dim=1)
if self.average == False:
yi = (yi == yi.max(dim=1,keepdim=True)[0]).float()
# deconv for reconsturction
wi_center = raw_wi[0]
yi = F.conv_transpose2d(yi, wi_center, stride=self.stride*self.scale, padding=self.scale)
yi =yi/6.
y.append(yi)
y = torch.cat(y, dim=0)
return y
4、Non-Local Sparse Attention
5、借鉴Asymmetric Non-local Neural Networks for Semantic Segmentation
https://blog.csdn.net/wildridder/article/details/104159458AFNB
import torch
from torch import nn
from torch.nn import functional as F
from models.tools.module_helper import ModuleHelper
class PSPModule(nn.Module):
# (1, 2, 3, 6)
def __init__(self, sizes=(1, 3, 6, 8), dimension=2):
super(PSPModule, self).__init__()
self.stages = nn.ModuleList([self._make_stage(size, dimension) for size in sizes])
def _make_stage(self, size, dimension=2):
if dimension == 1:
prior = nn.AdaptiveAvgPool1d(output_size=size)
elif dimension == 2:
prior = nn.AdaptiveAvgPool2d(output_size=(size, size))
elif dimension == 3:
prior = nn.AdaptiveAvgPool3d(output_size=(size, size, size))
return prior
def forward(self, feats):
n, c, _, _ = feats.size()
priors = [stage(feats).view(n, c, -1) for stage in self.stages]
center = torch.cat(priors, -1)
return center
class _SelfAttentionBlock(nn.Module):
'''
The basic implementation for self-attention block/non-local block
Input:
N X C X H X W
Parameters:
in_channels : the dimension of the input feature map
key_channels : the dimension after the key/query transform
value_channels : the dimension after the value transform
scale : choose the scale to downsample the input feature maps (save memory cost)
Return:
N X C X H X W
position-aware context features.(w/o concate or add with the input)
'''
def __init__(self, low_in_channels, high_in_channels, key_channels, value_channels, out_channels=None, scale=1, norm_type=None,psp_size=(1,3,6,8)):
super(_SelfAttentionBlock, self).__init__()
self.scale = scale
self.in_channels = low_in_channels
self.out_channels = out_channels
self.key_channels = key_channels
self.value_channels = value_channels
if out_channels == None:
self.out_channels = high_in_channels
self.pool = nn.MaxPool2d(kernel_size=(scale, scale))
self.f_key = nn.Sequential(
nn.Conv2d(in_channels=self.in_channels, out_channels=self.key_channels,
kernel_size=1, stride=1, padding=0),
ModuleHelper.BNReLU(self.key_channels, norm_type=norm_type),
)
self.f_query = nn.Sequential(
nn.Conv2d(in_channels=high_in_channels, out_channels=self.key_channels,
kernel_size=1, stride=1, padding=0),
ModuleHelper.BNReLU(self.key_channels, norm_type=norm_type),
)
self.f_value = nn.Conv2d(in_channels=self.in_channels, out_channels=self.value_channels,
kernel_size=1, stride=1, padding=0)
self.W = nn.Conv2d(in_channels=self.value_channels, out_channels=self.out_channels,
kernel_size=1, stride=1, padding=0)
self.psp = PSPModule(psp_size)
nn.init.constant_(self.W.weight, 0)
nn.init.constant_(self.W.bias, 0)
def forward(self, low_feats, high_feats):
batch_size, h, w = high_feats.size(0), high_feats.size(2), high_feats.size(3)
# if self.scale > 1:
# x = self.pool(x)
value = self.psp(self.f_value(low_feats))
query = self.f_query(high_feats).view(batch_size, self.key_channels, -1)
query = query.permute(0, 2, 1)
key = self.f_key(low_feats)
# value=self.psp(value)#.view(batch_size, self.value_channels, -1)
value = value.permute(0, 2, 1)
key = self.psp(key) # .view(batch_size, self.key_channels, -1)
sim_map = torch.matmul(query, key)
sim_map = (self.key_channels ** -.5) * sim_map
sim_map = F.softmax(sim_map, dim=-1)
context = torch.matmul(sim_map, value)
context = context.permute(0, 2, 1).contiguous()
context = context.view(batch_size, self.value_channels, *high_feats.size()[2:])
context = self.W(context)
return context
class SelfAttentionBlock2D(_SelfAttentionBlock):
def __init__(self, low_in_channels, high_in_channels, key_channels, value_channels, out_channels=None, scale=1, norm_type=None,psp_size=(1,3,6,8)):
super(SelfAttentionBlock2D, self).__init__(low_in_channels,
high_in_channels,
key_channels,
value_channels,
out_channels,
scale,
norm_type,
psp_size=psp_size
)
class AFNB(nn.Module):
"""
Parameters:
in_features / out_features: the channels of the input / output feature maps.
dropout: we choose 0.05 as the default value.
size: you can apply multiple sizes. Here we only use one size.
Return:
features fused with Object context information.
"""
def __init__(self, low_in_channels, high_in_channels, out_channels, key_channels, value_channels, dropout,
sizes=([1]), norm_type=None,psp_size=(1,3,6,8)):
super(AFNB, self).__init__()
self.stages = []
self.norm_type = norm_type
self.psp_size=psp_size
self.stages = nn.ModuleList(
[self._make_stage([low_in_channels, high_in_channels], out_channels, key_channels, value_channels, size) for
size in sizes])
self.conv_bn_dropout = nn.Sequential(
nn.Conv2d(out_channels + high_in_channels, out_channels, kernel_size=1, padding=0),
ModuleHelper.BatchNorm2d(norm_type=self.norm_type)(out_channels),
nn.Dropout2d(dropout)
)
def _make_stage(self, in_channels, output_channels, key_channels, value_channels, size):
return SelfAttentionBlock2D(in_channels[0],
in_channels[1],
key_channels,
value_channels,
output_channels,
size,
self.norm_type,
psp_size=self.psp_size)
def forward(self, low_feats, high_feats):
priors = [stage(low_feats, high_feats) for stage in self.stages]
context = priors[0]
for i in range(1, len(priors)):
context += priors[i]
output = self.conv_bn_dropout(torch.cat([context, high_feats], 1))
return output
APNB
'''
Modified from https://github.com/PkuRainBow/OCNet.pytorch/blob/master/oc_module/base_oc_block.py
'''
import torch
from torch import nn
from torch.nn import functional as F
from models.tools.module_helper import ModuleHelper
class PSPModule(nn.Module):
# (1, 2, 3, 6)
def __init__(self, sizes=(1, 3, 6, 8), dimension=2):
super(PSPModule, self).__init__()
self.stages = nn.ModuleList([self._make_stage(size, dimension) for size in sizes])
def _make_stage(self, size, dimension=2):
if dimension == 1:
prior = nn.AdaptiveAvgPool1d(output_size=size)
elif dimension == 2:
prior = nn.AdaptiveAvgPool2d(output_size=(size, size))
elif dimension == 3:
prior = nn.AdaptiveAvgPool3d(output_size=(size, size, size))
return prior
def forward(self, feats):
n, c, _, _ = feats.size()
priors = [stage(feats).view(n, c, -1) for stage in self.stages]
center = torch.cat(priors, -1)
return center
class _SelfAttentionBlock(nn.Module):
'''
The basic implementation for self-attention block/non-local block
Input:
N X C X H X W
Parameters:
in_channels : the dimension of the input feature map
key_channels : the dimension after the key/query transform
value_channels : the dimension after the value transform
scale : choose the scale to downsample the input feature maps (save memory cost)
Return:
N X C X H X W
position-aware context features.(w/o concate or add with the input)
'''
def __init__(self, in_channels, key_channels, value_channels, out_channels=None, scale=1, norm_type=None,psp_size=(1,3,6,8)):
super(_SelfAttentionBlock, self).__init__()
self.scale = scale
self.in_channels = in_channels
self.out_channels = out_channels
self.key_channels = key_channels
self.value_channels = value_channels
if out_channels == None:
self.out_channels = in_channels
self.pool = nn.MaxPool2d(kernel_size=(scale, scale))
self.f_key = nn.Sequential(
nn.Conv2d(in_channels=self.in_channels, out_channels=self.key_channels,
kernel_size=1, stride=1, padding=0),
ModuleHelper.BNReLU(self.key_channels, norm_type=norm_type),
)
self.f_query = self.f_key
self.f_value = nn.Conv2d(in_channels=self.in_channels, out_channels=self.value_channels,
kernel_size=1, stride=1, padding=0)
self.W = nn.Conv2d(in_channels=self.value_channels, out_channels=self.out_channels,
kernel_size=1, stride=1, padding=0)
self.psp = PSPModule(psp_size)
nn.init.constant_(self.W.weight, 0)
nn.init.constant_(self.W.bias, 0)
def forward(self, x):
batch_size, h, w = x.size(0), x.size(2), x.size(3)
if self.scale > 1:
x = self.pool(x)
value = self.psp(self.f_value(x))
query = self.f_query(x).view(batch_size, self.key_channels, -1)
query = query.permute(0, 2, 1)
key = self.f_key(x)
# value=self.psp(value)#.view(batch_size, self.value_channels, -1)
value = value.permute(0, 2, 1)
key = self.psp(key) # .view(batch_size, self.key_channels, -1)
sim_map = torch.matmul(query, key)
sim_map = (self.key_channels ** -.5) * sim_map
sim_map = F.softmax(sim_map, dim=-1)
context = torch.matmul(sim_map, value)
context = context.permute(0, 2, 1).contiguous()
context = context.view(batch_size, self.value_channels, *x.size()[2:])
context = self.W(context)
return context
class SelfAttentionBlock2D(_SelfAttentionBlock):
def __init__(self, in_channels, key_channels, value_channels, out_channels=None, scale=1, norm_type=None,psp_size=(1,3,6,8)):
super(SelfAttentionBlock2D, self).__init__(in_channels,
key_channels,
value_channels,
out_channels,
scale,
norm_type,
psp_size=psp_size)
class APNB(nn.Module):
"""
Parameters:
in_features / out_features: the channels of the input / output feature maps.
dropout: we choose 0.05 as the default value.
size: you can apply multiple sizes. Here we only use one size.
Return:
features fused with Object context information.
"""
def __init__(self, in_channels, out_channels, key_channels, value_channels, dropout, sizes=([1]), norm_type=None,psp_size=(1,3,6,8)):
super(APNB, self).__init__()
self.stages = []
self.norm_type = norm_type
self.psp_size=psp_size
self.stages = nn.ModuleList(
[self._make_stage(in_channels, out_channels, key_channels, value_channels, size) for size in sizes])
self.conv_bn_dropout = nn.Sequential(
nn.Conv2d(2 * in_channels, out_channels, kernel_size=1, padding=0),
ModuleHelper.BNReLU(out_channels, norm_type=norm_type),
nn.Dropout2d(dropout)
)
def _make_stage(self, in_channels, output_channels, key_channels, value_channels, size):
return SelfAttentionBlock2D(in_channels,
key_channels,
value_channels,
output_channels,
size,
self.norm_type,
self.psp_size)
def forward(self, feats):
priors = [stage(feats) for stage in self.stages]
context = priors[0]
for i in range(1, len(priors)):
context += priors[i]
output = self.conv_bn_dropout(torch.cat([context, feats], 1))
return output
简而言之，就是把H*W转化为了110，极大地精简了计算。
6559
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
