1 在 ultralytics/nn/attention 目录下创建MCA.py文件
import torch
from torch import nn
import math
__all__ = ['MCALayer', 'MCAGate']
class StdPool(nn.Module):
def __init__(self):
super(StdPool, self).__init__()
def forward(self, x):
b, c, _, _ = x.size()
std = x.view(b, c, -1).std(dim=2, keepdim=True)
std = std.reshape(b, c, 1, 1)
return std
class MCAGate(nn.Module):
def __init__(self, k_size, pool_types=['avg', 'std']):
"""Constructs a MCAGate module.
Args:
k_size: kernel size
pool_types: pooling type. 'avg': average pooling, 'max': max pooling, 'std': standard deviation pooling.
"""
super(MCAGate, self).__init__()
self.pools = nn.ModuleList([])
for pool_type in pool_types:
if pool_type == 'avg':
self.pools.append(nn.AdaptiveAvgPool2d(1))
elif pool_type == 'max':
self.pools.append(nn.AdaptiveMaxPool2d(1))
elif pool_type == 'std':
self.pools.append(StdPool())
else:
raise NotImplementedError
self.conv = nn.Conv2d(1, 1, kernel_size=(1, k_size), stride=1, padding=(0, (k_size - 1) // 2), bias=False)
self.sigmoid = nn.Sigmoid()
self.weight = nn.Parameter(torch.rand(2))
def forward(self, x):
feats = [pool(x) for pool in self.pools]
if len(feats) == 1:
out = feats[0]
elif len(feats) == 2:
weight = torch.sigmoid(self.weight)
out = 1 / 2 * (feats[0] + feats[1]) + weight[0] * feats[0] + weight[1] * feats[1]
else:
assert False, "Feature Extraction Exception!"
out = out.permute(0, 3, 2, 1).contiguous()
out = self.conv(out)
out = out.permute(0, 3, 2, 1).contiguous()
out = self.sigmoid(out)
out = out.expand_as(x)
return x * out
class MCALayer(nn.Module):
def __init__(self, inp, no_spatial=False):
"""Constructs a MCA module.
Args:
inp: Number of channels of the input feature maps
no_spatial: whether to build channel dimension interactions
"""
super(MCALayer, self).__init__()
lambd = 1.5
gamma = 1
temp = round(abs((math.log2(inp) - gamma) / lambd))
kernel = temp if temp % 2 else temp - 1
self.h_cw = MCAGate(3)
self.w_hc = MCAGate(3)
self.no_spatial = no_spatial
if not no_spatial:
self.c_hw = MCAGate(kernel)
def forward(self, x):
x_h = x.permute(0, 2, 1, 3).contiguous()
x_h = self.h_cw(x_h)
x_h = x_h.permute(0, 2, 1, 3).contiguous()
x_w = x.permute(0, 3, 2, 1).contiguous()
x_w = self.w_hc(x_w)
x_w = x_w.permute(0, 3, 2, 1).contiguous()
if not self.no_spatial:
x_c = self.c_hw(x)
x_out = 1 / 3 * (x_c + x_h + x_w)
else:
x_out = 1 / 2 * (x_h + x_w)
return x_out
2 修改ultralytics/nn/tasks.py文件
将下面代码添加到函数parse_model中
elif m in {MCALayer,MCAGate}:
args = [ch[f], *args]