1、RVBS
RVBS
SPD（Slicing）
python
class RVBS(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1):
super(RVBS, self).__init__()
out_channels = out_channels // 4
self.in_channels = in_channels
self.nonlinearity = nn.SiLU()
self.rbr_dense = self.conv_bn(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=padding)
self.rbr_1x1 = self.conv_bn(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=stride, padding=padding - kernel_size // 2)
self.rbr_identity = nn.BatchNorm2d(num_features=in_channels) if out_channels == in_channels and stride == 1 else None
def conv_bn(self, in_channels, out_channels, kernel_size, stride, padding):
result = nn.Sequential()
result.add_module('conv', nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=padding, bias=False))
result.add_module('bn', nn.BatchNorm2d(num_features=out_channels))
return result
def _pad_1x1_to_3x3_tensor(self, kernel1x1):
if kernel1x1 is None:
return 0
else:
return F.pad(kernel1x1, [1, 1, 1, 1])
def _fuse_bn_tensor(self, branch):
if branch is None:
return 0, 0
if isinstance(branch, nn.Sequential):
kernel = branch.conv.weight
running_mean = branch.bn.running_mean
running_var = branch.bn.running_var
gamma = branch.bn.weight
beta = branch.bn.bias
eps = branch.bn.eps
else:
assert isinstance(branch, nn.BatchNorm2d)
if not hasattr(self, 'id_tensor'):
kernel_value = np.zeros((self.in_channels, self.in_channels, 3, 3), dtype=np.float32)
for i in range(self.in_channels):
kernel_value[i, i, 1, 1] = 1
self.id_tensor = torch.from_numpy(kernel_value).to(branch.weight.device)
kernel = self.id_tensor
running_mean = branch.running_mean
running_var = branch.running_var
gamma = branch.weight
beta = branch.bias
eps = branch.eps
std = (running_var + eps).sqrt()
t = (gamma / std).reshape(-1, 1, 1, 1)
return kernel * t, beta - running_mean * gamma / std
def get_equivalent_kernel_bias(self):
kernel3x3, bias3x3 = self._fuse_bn_tensor(self.rbr_dense)
kernel1x1, bias1x1 = self._fuse_bn_tensor(self.rbr_1x1)
kernelid, biasid = self._fuse_bn_tensor(self.rbr_identity)
return kernel3x3 + self._pad_1x1_to_3x3_tensor(kernel1x1) + kernelid, bias3x3 + bias1x1 + biasid
def forward(self, inputs):
if self.rbr_identity is None:
id_out = 0
else:
id_out = self.rbr_identity(inputs)
x = self.nonlinearity(self.rbr_dense(inputs) + self.rbr_1x1(inputs) + id_out)
return torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1)
def fusevggforward(self, inputs):
x = self.nonlinearity(self.rbr_dense(inputs))
return torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1)2、ASDCN
ASDCN
python
class ASDCN(nn.Module):
def __init__(self, in_channels, out_channels, conv2d_bias=True, reduce_gamma=False):
super().__init__()
out_channels = out_channels // 4
self.square_conv = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1, bias=conv2d_bias)
self.square_bn = nn.BatchNorm2d(num_features=out_channels)
self.ver_conv = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=(3, 1), stride=1, padding=(1, 0), bias=conv2d_bias)
self.ver_bn = nn.BatchNorm2d(num_features=out_channels)
self.hor_conv = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=(1, 3), stride=1, padding=(0, 1), bias=conv2d_bias)
self.hor_bn = nn.BatchNorm2d(num_features=out_channels)
self.act = nn.SiLU()
if reduce_gamma:
self.init_gamma(1.0 / 3)
def init_gamma(self, gamma_value):
init.constant_(self.square_bn.weight, gamma_value)
init.constant_(self.ver_bn.weight, gamma_value)
init.constant_(self.hor_bn.weight, gamma_value)
def _fuse_bn_tensor(self, conv, bn):
std = (bn.running_var + bn.eps).sqrt()
t = (bn.weight / std).reshape(-1, 1, 1, 1)
conv_bias = conv.bias if conv.bias is not None else torch.zeros_like(bn.running_mean)
return conv.weight * t, bn.bias + (conv_bias - bn.running_mean) * bn.weight / std
def _add_to_square_kernel(self, square_kernel, asym_kernel):
asym_h = asym_kernel.size(2)
asym_w = asym_kernel.size(3)
square_h = square_kernel.size(2)
square_w = square_kernel.size(3)
square_kernel[:, :, square_h // 2 - asym_h // 2: square_h // 2 - asym_h // 2 + asym_h, square_w // 2 - asym_w // 2: square_w // 2 - asym_w // 2 + asym_w] += asym_kernel
def get_equivalent_kernel_bias(self):
hor_k, hor_b = self._fuse_bn_tensor(self.hor_conv, self.hor_bn)
ver_k, ver_b = self._fuse_bn_tensor(self.ver_conv, self.ver_bn)
square_k, square_b = self._fuse_bn_tensor(self.square_conv, self.square_bn)
self._add_to_square_kernel(square_k, hor_k)
self._add_to_square_kernel(square_k, ver_k)
return square_k, hor_b + ver_b + square_b
def fuse_convs(self):
deploy_k, deploy_b = self.get_equivalent_kernel_bias()
self.fused_conv = nn.Conv2d(in_channels=self.square_conv.in_channels,
out_channels=self.square_conv.out_channels,
kernel_size=self.square_conv.kernel_size,
stride=self.square_conv.stride,
padding=self.square_conv.padding,
bias=True).requires_grad_(False)
self.fused_conv.weight.data = deploy_k
self.fused_conv.bias.data = deploy_b
for para in self.parameters():
para.detach_()
self.__delattr__('square_conv')
self.__delattr__('square_bn')
self.__delattr__('hor_conv')
self.__delattr__('hor_bn')
self.__delattr__('ver_conv')
self.__delattr__('ver_bn')
def forward(self, input):
square_outputs = self.square_bn(self.square_conv(input))
vertical_outputs = self.ver_bn(self.ver_conv(input))
horizontal_outputs = self.hor_bn(self.hor_conv(input))
result = self.act(square_outputs + vertical_outputs + horizontal_outputs)
return torch.cat([result[..., ::2, ::2], result[..., 1::2, ::2], result[..., ::2, 1::2], result[..., 1::2, 1::2]], 1)
def forward_fuse(self, input):
result = self.act(self.fused_conv(input))
return torch.cat([result[..., ::2, ::2], result[..., 1::2, ::2], result[..., ::2, 1::2], result[..., 1::2, 1::2]], 1)3、 DBPA-C3k2
DBPA-C3k2
DBB
PPA
Attention
python
# --------------------------------------------------DBPA_C3k2 begin--------------------------------------------------
class DBPA_C3k2(nn.Module):
def __init__(self, c1, c2, n=1, c3k=False, e=0.5, front=False, middle=False, back=False, attention=None, single_init=True, shortcut=True):
super().__init__()
self.c = int(c2 * e)
self.cv1 = Conv(c1, 2 * self.c, 1, 1)
self.cv2 = Conv((2 + n) * self.c, c2, 1)
self.m = nn.ModuleList(nn.Sequential(ADBC3k(self.c, self.c, 2, middle, single_init, shortcut) if c3k else ADBBottleneck(self.c, self.c, middle, single_init, shortcut), DBB(self.c, self.c, single_init=single_init)) for _ in range(n))
self.front = front
self.back = back
if self.front:
# self.attention = LSKblock(c1)
# self.attention = CoordAtt(c1, c1)
# self.attention = SE(c1)
# self.attention = CBAM(c1)
# self.attention = OECA()
self.attention = PPA(c1, c1)
if self.back:
# if attention == 'LSKblock':
# self.attention = LSKblock(c2)
# if attention == 'CoordAtt':
# self.attention = CoordAtt(c2, c2)
# if attention == 'SE':
# self.attention = SE(c2)
# if attention == 'CBAM':
# self.attention = CBAM(c2)
# if attention == 'OECA':
# self.attention = OECA()
if attention is None:
self.attention = PPA(c2, c2)
def forward(self, x):
y = list(self.cv1(self.attention(x) if self.front else x).chunk(2, 1))
y.extend(m(y[-1]) for m in self.m)
return self.attention(self.cv2(torch.cat(y, 1))) if self.back else self.cv2(torch.cat(y, 1))
class ADBC3k(nn.Module):
def __init__(self, c1, c2, n=1, middle=False, single_init=True, shortcut=True, e=0.5):
super().__init__()
c_ = int(c2 * e)
self.cv1 = Conv(c1, c_, 1, 1)
self.cv2 = Conv(c1, c_, 1, 1)
self.cv3 = Conv(2 * c_, c2, 1)
self.m = nn.Sequential(*(ADBBottleneck(c_, c_, middle, single_init, shortcut, e=1.0) for _ in range(n)))
def forward(self, x):
return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), 1))
class ADBBottleneck(nn.Module):
def __init__(self, c1, c2, middle=False, single_init=True, shortcut=True, e=0.5):
super().__init__()
c_ = int(c2 * e)
self.cv1 = DBB(c1, c_, single_init=single_init)
self.cv2 = DBB(c_, c2, single_init=single_init)
self.middle = middle
if self.middle:
# self.attention = LSKblock(c2)
# self.attention = CoordAtt(c2, c2)
# self.attention = SE(c2)
# self.attention = CBAM(c2)
# self.attention = OECA()
self.attention = PPA(c2, c2)
self.add = shortcut and c1 == c2
def forward(self, x):
if self.middle:
return x + self.attention(self.cv2(self.cv1(x))) if self.add else self.attention(self.cv2(self.cv1(x)))
else:
return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))
# --------------------------------------------------DBPA_C3k2 end--------------------------------------------------
# ****************************************DBB begin****************************************
class DBB(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1, single_init=True):
super().__init__()
self.kernel_size = kernel_size
self.out_channels = out_channels
self.dbb_1x1 = conv_bn(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=stride, padding=0)
self.dbb_1x1_kxk = nn.Sequential()
self.dbb_1x1_kxk.add_module('idconv1', IdentityBasedConv1x1(channels=in_channels))
self.dbb_1x1_kxk.add_module('bn1', BNAndPadLayer(pad_pixels=padding, num_features=in_channels))
self.dbb_1x1_kxk.add_module('conv2', nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=0, bias=False))
self.dbb_1x1_kxk.add_module('bn2', nn.BatchNorm2d(out_channels))
self.dbb_avg = nn.Sequential()
self.dbb_avg.add_module('conv', nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=1, padding=0, bias=False))
self.dbb_avg.add_module('bn', BNAndPadLayer(pad_pixels=padding, num_features=out_channels))
self.dbb_avg.add_module('avg', nn.AvgPool2d(kernel_size=kernel_size, stride=stride, padding=0))
self.dbb_avg.add_module('avgbn', nn.BatchNorm2d(out_channels))
self.dbb_origin = conv_bn(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=padding)
self.nonlinear = nn.SiLU()
if single_init:
self.single_init()
def single_init(self):
self.init_gamma(0.0)
init.constant_(self.dbb_origin.bn.weight, 1.0)
def init_gamma(self, gamma_value):
init.constant_(self.dbb_origin.bn.weight, gamma_value)
init.constant_(self.dbb_1x1.bn.weight, gamma_value)
init.constant_(self.dbb_avg.avgbn.weight, gamma_value)
init.constant_(self.dbb_1x1_kxk.bn2.weight, gamma_value)
def get_equivalent_kernel_bias(self):
k_origin, b_origin = transI_fusebn(self.dbb_origin.conv.weight, self.dbb_origin.bn)
k_1x1, b_1x1 = transI_fusebn(self.dbb_1x1.conv.weight, self.dbb_1x1.bn)
k_1x1 = transVI_multiscale(k_1x1, self.kernel_size)
k_1x1_kxk_first = self.dbb_1x1_kxk.idconv1.get_actual_kernel()
k_1x1_kxk_first, b_1x1_kxk_first = transI_fusebn(k_1x1_kxk_first, self.dbb_1x1_kxk.bn1)
k_1x1_kxk_second, b_1x1_kxk_second = transI_fusebn(self.dbb_1x1_kxk.conv2.weight, self.dbb_1x1_kxk.bn2)
k_1x1_kxk_merged, b_1x1_kxk_merged = transIII_1x1_kxk(k_1x1_kxk_first, b_1x1_kxk_first, k_1x1_kxk_second, b_1x1_kxk_second, groups=1)
k_avg = transV_avg(self.out_channels, self.kernel_size, 1)
k_1x1_avg_second, b_1x1_avg_second = transI_fusebn(k_avg.to(self.dbb_avg.avgbn.weight.device), self.dbb_avg.avgbn)
k_1x1_avg_first, b_1x1_avg_first = transI_fusebn(self.dbb_avg.conv.weight, self.dbb_avg.bn)
k_1x1_avg_merged, b_1x1_avg_merged = transIII_1x1_kxk(k_1x1_avg_first, b_1x1_avg_first, k_1x1_avg_second, b_1x1_avg_second, groups=1)
return transII_addbranch((k_origin, k_1x1, k_1x1_kxk_merged, k_1x1_avg_merged), (b_origin, b_1x1, b_1x1_kxk_merged, b_1x1_avg_merged))
def fuse_convs(self):
kernel, bias = self.get_equivalent_kernel_bias()
self.dbb_reparam = nn.Conv2d(in_channels=self.dbb_origin.conv.in_channels,
out_channels=self.dbb_origin.conv.out_channels,
kernel_size=self.dbb_origin.conv.kernel_size,
stride=self.dbb_origin.conv.stride,
padding=self.dbb_origin.conv.padding,
bias=True).requires_grad_(False)
self.dbb_reparam.weight.data = kernel
self.dbb_reparam.bias.data = bias
for para in self.parameters():
para.detach_()
self.__delattr__('dbb_origin')
self.__delattr__('dbb_avg')
self.__delattr__('dbb_1x1')
self.__delattr__('dbb_1x1_kxk')
def forward(self, inputs):
return self.nonlinear(self.dbb_origin(inputs) + self.dbb_1x1(inputs) + self.dbb_avg(inputs) + self.dbb_1x1_kxk(inputs))
def forward_fuse(self, inputs):
return self.nonlinear(self.dbb_reparam(inputs))
def conv_bn(in_channels, out_channels, kernel_size, stride, padding):
conv_layer = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=padding, bias=False)
bn_layer = nn.BatchNorm2d(num_features=out_channels)
se = nn.Sequential()
se.add_module('conv', conv_layer)
se.add_module('bn', bn_layer)
return se
class IdentityBasedConv1x1(nn.Conv2d):
def __init__(self, channels):
super().__init__(in_channels=channels, out_channels=channels, kernel_size=1, stride=1, padding=0, bias=False)
id_value = np.zeros((channels, channels, 1, 1))
for i in range(channels):
id_value[i, i, 0, 0] = 1
self.id_tensor = torch.from_numpy(id_value).type_as(self.weight)
init.zeros_(self.weight)
def forward(self, input):
kernel = self.weight + self.id_tensor.to(self.weight.device)
if input.dtype != kernel.dtype:
kernel = kernel.to(input.dtype)
result = F.conv2d(input, kernel, None, stride=1, padding=0, dilation=self.dilation, groups=self.groups)
return result
def get_actual_kernel(self):
return self.weight + self.id_tensor.to(self.weight.device)
class BNAndPadLayer(nn.Module):
def __init__(self, pad_pixels, num_features):
super().__init__()
self.bn = nn.BatchNorm2d(num_features)
self.pad_pixels = pad_pixels
def forward(self, input):
output = self.bn(input)
if self.pad_pixels > 0:
pad_values = self.bn.bias.detach() - self.bn.running_mean * self.bn.weight.detach() / torch.sqrt(self.bn.running_var + self.bn.eps)
output = F.pad(output, [self.pad_pixels] * 4)
pad_values = pad_values.view(1, -1, 1, 1)
output[:, :, 0:self.pad_pixels, :] = pad_values
output[:, :, -self.pad_pixels:, :] = pad_values
output[:, :, :, 0:self.pad_pixels] = pad_values
output[:, :, :, -self.pad_pixels:] = pad_values
return output
@property
def weight(self):
return self.bn.weight
@property
def bias(self):
return self.bn.bias
@property
def running_mean(self):
return self.bn.running_mean
@property
def running_var(self):
return self.bn.running_var
@property
def eps(self):
return self.bn.eps
def transI_fusebn(kernel, bn):
gamma = bn.weight
std = (bn.running_var + bn.eps).sqrt()
return kernel * ((gamma / std).reshape(-1, 1, 1, 1)), bn.bias - bn.running_mean * gamma / std
def transII_addbranch(kernels, biases):
return sum(kernels), sum(biases)
def transIII_1x1_kxk(k1, b1, k2, b2, groups):
if groups == 1:
k = F.conv2d(k2, k1.permute(1, 0, 2, 3))
b_hat = (k2 * b1.reshape(1, -1, 1, 1)).sum((1, 2, 3))
else:
k_slices = []
b_slices = []
k1_T = k1.permute(1, 0, 2, 3)
k1_group_width = k1.size(0) // groups
k2_group_width = k2.size(0) // groups
for g in range(groups):
k1_T_slice = k1_T[:, g * k1_group_width:(g + 1) * k1_group_width, :, :]
k2_slice = k2[g * k2_group_width:(g + 1) * k2_group_width, :, :, :]
k_slices.append(F.conv2d(k2_slice, k1_T_slice))
b_slices.append((k2_slice * b1[g * k1_group_width:(g + 1) * k1_group_width].reshape(1, -1, 1, 1)).sum((1, 2, 3)))
k, b_hat = transIV_depthconcat(k_slices, b_slices)
return k, b_hat + b2
def transIV_depthconcat(kernels, biases):
return torch.cat(kernels, dim=0), torch.cat(biases)
def transV_avg(channels, kernel_size, groups):
input_dim = channels // groups
k = torch.zeros((channels, input_dim, kernel_size, kernel_size))
k[np.arange(channels), np.tile(np.arange(input_dim), groups), :, :] = 1.0 / kernel_size ** 2
return k
def transVI_multiscale(kernel, target_kernel_size):
H_pixels_to_pad = (target_kernel_size - kernel.size(2)) // 2
W_pixels_to_pad = (target_kernel_size - kernel.size(3)) // 2
return F.pad(kernel, [H_pixels_to_pad, H_pixels_to_pad, W_pixels_to_pad, W_pixels_to_pad])
# ****************************************DBB end******************************************
# ****************************************PPA begin****************************************
class SpatialAttentionModule(nn.Module):
def __init__(self):
super(SpatialAttentionModule, self).__init__()
self.conv2d = nn.Conv2d(
in_channels=2, out_channels=1, kernel_size=7, stride=1, padding=3
)
self.sigmoid = nn.Sigmoid()
def forward(self, x):
avgout = torch.mean(x, dim=1, keepdim=True)
maxout, _ = torch.max(x, dim=1, keepdim=True)
out = torch.cat([avgout, maxout], dim=1)
out = self.sigmoid(self.conv2d(out))
return out * x
class LocalGlobalAttention(nn.Module):
def __init__(self, output_dim, patch_size):
super().__init__()
self.output_dim = output_dim
self.patch_size = patch_size
self.mlp1 = nn.Linear(patch_size * patch_size, output_dim // 2)
self.norm = nn.LayerNorm(output_dim // 2)
self.mlp2 = nn.Linear(output_dim // 2, output_dim)
self.conv = nn.Conv2d(output_dim, output_dim, kernel_size=1)
self.prompt = torch.nn.parameter.Parameter(
torch.randn(output_dim, requires_grad=True)
)
self.top_down_transform = torch.nn.parameter.Parameter(
torch.eye(output_dim), requires_grad=True
)
def forward(self, x):
x = x.permute(0, 2, 3, 1)
B, H, W, C = x.shape
P = self.patch_size
# Local branch
local_patches = x.unfold(1, P, P).unfold(2, P, P) # (B, H/P, W/P, P, P, C)
local_patches = local_patches.reshape(B, -1, P * P, C) # (B, H/P*W/P, P*P, C)
local_patches = local_patches.mean(dim=-1) # (B, H/P*W/P, P*P)
local_patches = self.mlp1(local_patches) # (B, H/P*W/P, input_dim // 2)
local_patches = self.norm(local_patches) # (B, H/P*W/P, input_dim // 2)
local_patches = self.mlp2(local_patches) # (B, H/P*W/P, output_dim)
local_attention = F.softmax(local_patches, dim=-1) # (B, H/P*W/P, output_dim)
local_out = local_patches * local_attention # (B, H/P*W/P, output_dim)
cos_sim = F.normalize(local_out, dim=-1) @ F.normalize(
self.prompt[None, ..., None], dim=1
) # B, N, 1
mask = cos_sim.clamp(0, 1) # B, N, 1
local_out = local_out * mask # (B, H/P*W/P, output_dim)
local_out = local_out @ self.top_down_transform # (B, H/P*W/P, output_dim)
# Restore shapes
local_out = local_out.reshape(
B, H // P, W // P, self.output_dim
) # (B, H/P, W/P, output_dim)
local_out = local_out.permute(0, 3, 1, 2)
local_out = F.interpolate(
local_out, size=(H, W), mode="bilinear", align_corners=False
) # (B, output_dim,H, W, )
output = self.conv(local_out) # (B, output_dim,H, W, )
return output
class ECA(nn.Module):
def __init__(self, in_channel, gamma=2, b=1):
super(ECA, self).__init__()
k = int(abs((math.log(in_channel, 2) + b) / gamma))
kernel_size = k if k % 2 else k + 1
padding = kernel_size // 2
self.pool = nn.AdaptiveAvgPool2d(output_size=1)
self.conv = nn.Sequential(
nn.Conv1d(
in_channels=1,
out_channels=1,
kernel_size=kernel_size,
padding=padding,
bias=False,
),
nn.Sigmoid(),
)
def forward(self, x):
out = self.pool(x)
out = out.view(x.size(0), 1, x.size(1))
out = self.conv(out)
out = out.view(x.size(0), x.size(1), 1, 1)
return out * x
class conv_block(nn.Module):
def __init__(
self,
in_features,
out_features,
kernel_size=(3, 3),
stride=(1, 1),
padding=(1, 1),
dilation=(1, 1),
norm_type="bn",
activation=True,
use_bias=True,
groups=1,
):
super().__init__()
self.conv = nn.Conv2d(
in_channels=in_features,
out_channels=out_features,
kernel_size=kernel_size,
stride=stride,
padding=padding,
dilation=dilation,
bias=use_bias,
groups=groups,
)
self.norm_type = norm_type
self.act = activation
if self.norm_type == "gn":
self.norm = nn.GroupNorm(
32 if out_features >= 32 else out_features, out_features
)
if self.norm_type == "bn":
self.norm = nn.BatchNorm2d(out_features)
if self.act:
# self.relu = nn.GELU()
self.relu = nn.ReLU(inplace=False)
def forward(self, x):
x = self.conv(x)
if self.norm_type is not None:
x = self.norm(x)
if self.act:
x = self.relu(x)
return x
class PPA(nn.Module):
def __init__(self, in_features, filters) -> None:
super().__init__()
self.skip = conv_block(
in_features=in_features,
out_features=filters,
kernel_size=(1, 1),
padding=(0, 0),
norm_type="bn",
activation=False,
)
self.c1 = conv_block(
in_features=in_features,
out_features=filters,
kernel_size=(3, 3),
padding=(1, 1),
norm_type="bn",
activation=True,
)
self.c2 = conv_block(
in_features=filters,
out_features=filters,
kernel_size=(3, 3),
padding=(1, 1),
norm_type="bn",
activation=True,
)
self.c3 = conv_block(
in_features=filters,
out_features=filters,
kernel_size=(3, 3),
padding=(1, 1),
norm_type="bn",
activation=True,
)
self.sa = SpatialAttentionModule()
self.cn = ECA(filters)
self.lga2 = LocalGlobalAttention(filters, 2)
self.lga4 = LocalGlobalAttention(filters, 4)
self.bn1 = nn.BatchNorm2d(filters)
self.drop = nn.Dropout2d(0.1)
self.relu = nn.ReLU()
self.gelu = nn.GELU()
def forward(self, x):
x_skip = self.skip(x)
x_lga2 = self.lga2(x_skip)
x_lga4 = self.lga4(x_skip)
x1 = self.c1(x)
x2 = self.c2(x1)
x3 = self.c3(x2)
x = x1 + x2 + x3 + x_skip + x_lga2 + x_lga4
x = self.cn(x)
x = self.sa(x)
x = self.drop(x)
x = self.bn1(x)
x = self.relu(x)
return x
# ****************************************PPA****************************************4、 FI-WIoUv2
python
class WIoU_Scale1:
''' if monotonous = None , v1
if monotonous = True , v2
if monotonous = False , v3
'''
iou_mean = 1.
monotonous = None
_momentum = 1 - 0.5 ** (1 / 7000)
_is_train = True
def __init__(self, iou):
self.iou = iou
self._update(self)
@classmethod
def _update(cls, self):
if cls._is_train: cls.iou_mean = (1 - cls._momentum) * cls.iou_mean + cls._momentum * self.iou.detach().mean().item()
@classmethod
def _scaled_loss(cls, self, gamma=1.9, delta=3):
if isinstance(self.monotonous, bool):
if self.monotonous:
return (self.iou.detach() / self.iou_mean).sqrt()
else:
beta = self.iou.detach() / self.iou_mean
alpha = delta * torch.pow(gamma, beta - delta)
return beta / alpha
return 1
class WIoU_Scale2:
''' if monotonous = None , v1
if monotonous = True , v2
if monotonous = False , v3
'''
iou_mean = 1.
monotonous = True
_momentum = 1 - 0.5 ** (1 / 7000)
_is_train = True
def __init__(self, iou):
self.iou = iou
self._update(self)
@classmethod
def _update(cls, self):
if cls._is_train: cls.iou_mean = (1 - cls._momentum) * cls.iou_mean + \
cls._momentum * self.iou.detach().mean().item()
@classmethod
def _scaled_loss(cls, self, gamma=1.9, delta=3):
if isinstance(self.monotonous, bool):
if self.monotonous:
return (self.iou.detach() / self.iou_mean).sqrt()
else:
beta = self.iou.detach() / self.iou_mean
alpha = delta * torch.pow(gamma, beta - delta)
return beta / alpha
return 1
class WIoU_Scale3:
''' if monotonous = None , v1
if monotonous = True , v2
if monotonous = False , v3
'''
iou_mean = 1.
monotonous = False
_momentum = 1 - 0.5 ** (1 / 7000)
_is_train = True
def __init__(self, iou):
self.iou = iou
self._update(self)
@classmethod
def _update(cls, self):
if cls._is_train: cls.iou_mean = (1 - cls._momentum) * cls.iou_mean + \
cls._momentum * self.iou.detach().mean().item()
@classmethod
def _scaled_loss(cls, self, gamma=1.9, delta=3):
if isinstance(self.monotonous, bool):
if self.monotonous:
return (self.iou.detach() / self.iou_mean).sqrt()
else:
beta = self.iou.detach() / self.iou_mean
alpha = delta * torch.pow(gamma, beta - delta)
return beta / alpha
return 1
def bbox_iou(box1, box2, xywh=True, GIoU=False, DIoU=False, CIoU=False, EIoU=False, SIoU=False, MPDIoU1=False, MPDIoU2=False, MPDIoU3=False, WIoU1=False, WIoU2=False, WIoU3=False, scale=False, Focaler=False, Focal=False, Inner=False, ratio=1.0, d=0.00, u=0.95, eps=1e-7):
"""
Calculate Intersection over Union (IoU) of box1(1, 4) to box2(n, 4).
Args:
box1 (torch.Tensor): A tensor representing a single bounding box with shape (1, 4).
box2 (torch.Tensor): A tensor representing n bounding boxes with shape (n, 4).
xywh (bool, optional): If True, input boxes are in (x, y, w, h) format. If False, input boxes are in
(x1, y1, x2, y2) format. Defaults to True.
GIoU (bool, optional): If True, calculate Generalized IoU. Defaults to False.
DIoU (bool, optional): If True, calculate Distance IoU. Defaults to False.
CIoU (bool, optional): If True, calculate Complete IoU. Defaults to False.
eps (float, optional): A small value to avoid division by zero. Defaults to 1e-7.
Returns:
(torch.Tensor): IoU, GIoU, DIoU, or CIoU values depending on the specified flags.
"""
# Get the coordinates of bounding boxes
if xywh: # transform from xywh to xyxy
(x1, y1, w1, h1), (x2, y2, w2, h2) = box1.chunk(4, -1), box2.chunk(4, -1)
w1_, h1_, w2_, h2_ = w1 / 2, h1 / 2, w2 / 2, h2 / 2
b1_x1, b1_x2, b1_y1, b1_y2 = x1 - w1_, x1 + w1_, y1 - h1_, y1 + h1_
b2_x1, b2_x2, b2_y1, b2_y2 = x2 - w2_, x2 + w2_, y2 - h2_, y2 + h2_
else: # x1, y1, x2, y2 = box1
b1_x1, b1_y1, b1_x2, b1_y2 = box1.chunk(4, -1)
b2_x1, b2_y1, b2_x2, b2_y2 = box2.chunk(4, -1)
w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1 + eps
w2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1 + eps
if Inner:
x1 = (b1_x1 + b1_x2) / 2
y1 = (b1_y1 + b1_y2) / 2
x2 = (b2_x1 + b2_x2) / 2
y2 = (b2_y1 + b2_y2) / 2
w1_, h1_, w2_, h2_ = w1 / 2, h1 / 2, w2 / 2, h2 / 2
inner_b1_x1, inner_b1_x2, inner_b1_y1, inner_b1_y2 = x1 - w1_ * ratio, x1 + w1_ * ratio, y1 - h1_ * ratio, y1 + h1_ * ratio
inner_b2_x1, inner_b2_x2, inner_b2_y1, inner_b2_y2 = x2 - w2_ * ratio, x2 + w2_ * ratio, y2 - h2_ * ratio, y2 + h2_ * ratio
inner_inter = (torch.min(inner_b1_x2, inner_b2_x2) - torch.max(inner_b1_x1, inner_b2_x1)).clamp(0) * (torch.min(inner_b1_y2, inner_b2_y2) - torch.max(inner_b1_y1, inner_b2_y1)).clamp(0)
inner_union = w1 * ratio * h1 * ratio + w2 * ratio * h2 * ratio - inner_inter + eps
inner_iou = inner_inter / inner_union
# Intersection area
inter = (b1_x2.minimum(b2_x2) - b1_x1.maximum(b2_x1)).clamp_(0) * (b1_y2.minimum(b2_y2) - b1_y1.maximum(b2_y1)).clamp_(0)
# Union Area
union = w1 * h1 + w2 * h2 - inter + eps
# IoU
iou = inter / union
if scale:
if Focaler:
if Inner:
iou = ((inner_iou - d) / (u - d)).clamp(0, 1) # default d=0.00,u=0.95
else:
iou = ((iou - d) / (u - d)).clamp(0, 1) # default d=0.00,u=0.95
if WIoU1:
self = WIoU_Scale1(1 - iou)
elif WIoU2:
self = WIoU_Scale2(1 - iou)
else:
self = WIoU_Scale3(1 - iou)
if CIoU or DIoU or GIoU or EIoU or SIoU or WIoU1 or WIoU2 or WIoU3 or MPDIoU1 or MPDIoU2 or MPDIoU3:
cw = b1_x2.maximum(b2_x2) - b1_x1.minimum(b2_x1) # convex (smallest enclosing box) width
ch = b1_y2.maximum(b2_y2) - b1_y1.minimum(b2_y1) # convex height
if CIoU or DIoU or EIoU or SIoU or WIoU1 or WIoU2 or WIoU3 or MPDIoU1 or MPDIoU2 or MPDIoU3: # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1
c2 = cw.pow(2) + ch.pow(2) + eps # convex diagonal squared
rho2 = (
(b2_x1 + b2_x2 - b1_x1 - b1_x2).pow(2) + (b2_y1 + b2_y2 - b1_y1 - b1_y2).pow(2)
) / 4 # center dist**2
if CIoU: # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47
v = (4 / math.pi ** 2) * ((w2 / h2).atan() - (w1 / h1).atan()).pow(2)
with torch.no_grad():
alpha = v / (v - iou + (1 + eps))
if Focaler:
if Inner:
iou = ((inner_iou - d) / (u - d)).clamp(0, 1) # default d=0.00,u=0.95
else:
iou = ((iou - d) / (u - d)).clamp(0, 1) # default d=0.00,u=0.95
if Focal:
return iou - (rho2 / c2 + v * alpha), torch.pow(iou, 0.5)
else:
return iou - (rho2 / c2 + v * alpha) # CIoU
if EIoU:
rho_w2 = ((b2_x2 - b2_x1) - (b1_x2 - b1_x1)).pow(2)
rho_h2 = ((b2_y2 - b2_y1) - (b1_y2 - b1_y1)).pow(2)
cw2 = cw.pow(2) + eps
ch2 = ch.pow(2) + eps
if Focaler:
iou = ((iou - d) / (u - d)).clamp(0, 1)
if Focal:
return iou - (rho2 / c2 + rho_w2 / cw2 + rho_h2 / ch2), torch.pow(iou, 0.5)
else:
return iou - (rho2 / c2 + rho_w2 / cw2 + rho_h2 / ch2)
if SIoU:
s_cw = (b2_x1 + b2_x2 - b1_x1 - b1_x2) * 0.5 + eps
s_ch = (b2_y1 + b2_y2 - b1_y1 - b1_y2) * 0.5 + eps
sigma = torch.pow(s_cw ** 2 + s_ch ** 2, 0.5)
sin_alpha_1 = torch.abs(s_cw) / sigma
sin_alpha_2 = torch.abs(s_ch) / sigma
threshold = pow(2, 0.5) / 2
sin_alpha = torch.where(sin_alpha_1 > threshold, sin_alpha_2, sin_alpha_1)
angle_cost = torch.cos(torch.arcsin(sin_alpha) * 2 - math.pi / 2)
rho_x = (s_cw / cw) ** 2
rho_y = (s_ch / ch) ** 2
gamma = angle_cost - 2
distance_cost = 2 - torch.exp(gamma * rho_x) - torch.exp(gamma * rho_y)
omiga_w = torch.abs(w1 - w2) / torch.max(w1, w2)
omiga_h = torch.abs(h1 - h2) / torch.max(h1, h2)
shape_cost = torch.pow(1 - torch.exp(-1 * omiga_w), 4) + torch.pow(1 - torch.exp(-1 * omiga_h), 4)
if Focaler:
iou = ((iou - d) / (u - d)).clamp(0, 1)
if Focal:
return iou - 0.5 * (distance_cost + shape_cost), torch.pow(iou, 0.5)
else:
return iou - 0.5 * (distance_cost + shape_cost)
if WIoU1:
if scale:
return getattr(WIoU_Scale1, '_scaled_loss')(self), (1 - iou) * torch.exp((rho2 / c2)), iou # WIoU https://arxiv.org/abs/2301.10051
else:
return iou, torch.exp((rho2 / c2)) # WIoU v1
if WIoU2:
if scale:
return getattr(WIoU_Scale2, '_scaled_loss')(self), (1 - iou) * torch.exp((rho2 / c2)), iou # WIoU https://arxiv.org/abs/2301.10051
else:
return iou, torch.exp((rho2 / c2)) # WIoU v2
if WIoU3:
if scale:
return getattr(WIoU_Scale3, '_scaled_loss')(self), (1 - iou) * torch.exp((rho2 / c2)), iou # WIoU https://arxiv.org/abs/2301.10051
else:
return iou, torch.exp((rho2 / c2)) # WIoU v3
if MPDIoU1:
cw2 = cw.pow(2) + eps
ch2 = ch.pow(2) + eps
d12 = ((b2_x1 - b1_x1) - (b2_y1 - b1_y1)) ** 2
d22 = ((b2_x2 - b1_x2) - (b2_y2 - b1_y2)) ** 2
if Focaler:
iou = ((iou - d) / (u - d)).clamp(0, 1)
if Focal:
return iou - ((d12 + d22) / (cw2 + ch2)), torch.pow(iou, 0.5)
else:
return iou - ((d12 + d22) / (cw2 + ch2))
if MPDIoU2:
cw2 = cw.pow(2)
ch2 = ch.pow(2)
d12 = (b2_x1 - b1_x1) ** 2 + (b2_y1 - b1_y1) ** 2
d22 = (b2_x2 - b1_x2) ** 2 + (b2_y2 - b1_y2) ** 2
if Focaler:
iou = ((iou - d) / (u - d)).clamp(0, 1)
if Focal:
return iou - ((d12 + d22) / (cw2 + ch2 + eps)), torch.pow(iou, 0.5)
else:
return iou - ((d12 + d22) / (cw2 + ch2 + eps))
if MPDIoU3:
sq_sum = (640 ** 2) + (640 ** 2)
d12 = (b2_x1 - b1_x1) ** 2 + (b2_y1 - b1_y1) ** 2
d22 = (b2_x2 - b1_x2) ** 2 + (b2_y2 - b1_y2) ** 2
if Focaler:
iou = ((iou - d) / (u - d)).clamp(0, 1)
if Focal:
return iou - ((d12 + d22) / sq_sum), torch.pow(iou, 0.5)
else:
return iou - ((d12 + d22) / sq_sum)
if Focaler:
iou = ((iou - d) / (u - d)).clamp(0, 1)
if Focal:
return iou - rho2 / c2, torch.pow(iou, 0.5)
else:
return iou - rho2 / c2 # DIoU
c_area = cw * ch + eps # convex area
if Focaler:
iou = ((iou - d) / (u - d)).clamp(0, 1)
if Focal:
return iou - (c_area - union) / c_area, torch.pow(iou, 0.5)
else:
return iou - (c_area - union) / c_area # GIoU https://arxiv.org/pdf/1902.09630.pdf
return iou # IoU