|
""" |
|
Feature Fusion for Varible-Length Data Processing |
|
AFF/iAFF is referred and modified from https://github.com/YimianDai/open-aff/blob/master/aff_pytorch/aff_net/fusion.py |
|
According to the paper: Yimian Dai et al, Attentional Feature Fusion, IEEE Winter Conference on Applications of Computer Vision, WACV 2021 |
|
""" |
|
|
|
import torch |
|
import torch.nn as nn |
|
|
|
|
|
class DAF(nn.Module): |
|
""" |
|
直接相加 DirectAddFuse |
|
""" |
|
|
|
def __init__(self): |
|
super(DAF, self).__init__() |
|
|
|
def forward(self, x, residual): |
|
return x + residual |
|
|
|
|
|
class iAFF(nn.Module): |
|
""" |
|
多特征融合 iAFF |
|
""" |
|
|
|
def __init__(self, channels=64, r=4, type="2D"): |
|
super(iAFF, self).__init__() |
|
inter_channels = int(channels // r) |
|
|
|
if type == "1D": |
|
|
|
self.local_att = nn.Sequential( |
|
nn.Conv1d(channels, inter_channels, kernel_size=1, stride=1, padding=0), |
|
nn.BatchNorm1d(inter_channels), |
|
nn.ReLU(inplace=True), |
|
nn.Conv1d(inter_channels, channels, kernel_size=1, stride=1, padding=0), |
|
nn.BatchNorm1d(channels), |
|
) |
|
|
|
|
|
self.global_att = nn.Sequential( |
|
nn.AdaptiveAvgPool1d(1), |
|
nn.Conv1d(channels, inter_channels, kernel_size=1, stride=1, padding=0), |
|
nn.BatchNorm1d(inter_channels), |
|
nn.ReLU(inplace=True), |
|
nn.Conv1d(inter_channels, channels, kernel_size=1, stride=1, padding=0), |
|
nn.BatchNorm1d(channels), |
|
) |
|
|
|
|
|
self.local_att2 = nn.Sequential( |
|
nn.Conv1d(channels, inter_channels, kernel_size=1, stride=1, padding=0), |
|
nn.BatchNorm1d(inter_channels), |
|
nn.ReLU(inplace=True), |
|
nn.Conv1d(inter_channels, channels, kernel_size=1, stride=1, padding=0), |
|
nn.BatchNorm1d(channels), |
|
) |
|
|
|
self.global_att2 = nn.Sequential( |
|
nn.AdaptiveAvgPool1d(1), |
|
nn.Conv1d(channels, inter_channels, kernel_size=1, stride=1, padding=0), |
|
nn.BatchNorm1d(inter_channels), |
|
nn.ReLU(inplace=True), |
|
nn.Conv1d(inter_channels, channels, kernel_size=1, stride=1, padding=0), |
|
nn.BatchNorm1d(channels), |
|
) |
|
elif type == "2D": |
|
|
|
self.local_att = nn.Sequential( |
|
nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0), |
|
nn.BatchNorm2d(inter_channels), |
|
nn.ReLU(inplace=True), |
|
nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0), |
|
nn.BatchNorm2d(channels), |
|
) |
|
|
|
|
|
self.global_att = nn.Sequential( |
|
nn.AdaptiveAvgPool2d(1), |
|
nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0), |
|
nn.BatchNorm2d(inter_channels), |
|
nn.ReLU(inplace=True), |
|
nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0), |
|
nn.BatchNorm2d(channels), |
|
) |
|
|
|
|
|
self.local_att2 = nn.Sequential( |
|
nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0), |
|
nn.BatchNorm2d(inter_channels), |
|
nn.ReLU(inplace=True), |
|
nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0), |
|
nn.BatchNorm2d(channels), |
|
) |
|
|
|
self.global_att2 = nn.Sequential( |
|
nn.AdaptiveAvgPool2d(1), |
|
nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0), |
|
nn.BatchNorm2d(inter_channels), |
|
nn.ReLU(inplace=True), |
|
nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0), |
|
nn.BatchNorm2d(channels), |
|
) |
|
else: |
|
raise f"the type is not supported" |
|
|
|
self.sigmoid = nn.Sigmoid() |
|
|
|
def forward(self, x, residual): |
|
flag = False |
|
xa = x + residual |
|
if xa.size(0) == 1: |
|
xa = torch.cat([xa, xa], dim=0) |
|
flag = True |
|
xl = self.local_att(xa) |
|
xg = self.global_att(xa) |
|
xlg = xl + xg |
|
wei = self.sigmoid(xlg) |
|
xi = x * wei + residual * (1 - wei) |
|
|
|
xl2 = self.local_att2(xi) |
|
xg2 = self.global_att(xi) |
|
xlg2 = xl2 + xg2 |
|
wei2 = self.sigmoid(xlg2) |
|
xo = x * wei2 + residual * (1 - wei2) |
|
if flag: |
|
xo = xo[0].unsqueeze(0) |
|
return xo |
|
|
|
|
|
class AFF(nn.Module): |
|
""" |
|
多特征融合 AFF |
|
""" |
|
|
|
def __init__(self, channels=64, r=4, type="2D"): |
|
super(AFF, self).__init__() |
|
inter_channels = int(channels // r) |
|
|
|
if type == "1D": |
|
self.local_att = nn.Sequential( |
|
nn.Conv1d(channels, inter_channels, kernel_size=1, stride=1, padding=0), |
|
nn.BatchNorm1d(inter_channels), |
|
nn.ReLU(inplace=True), |
|
nn.Conv1d(inter_channels, channels, kernel_size=1, stride=1, padding=0), |
|
nn.BatchNorm1d(channels), |
|
) |
|
self.global_att = nn.Sequential( |
|
nn.AdaptiveAvgPool1d(1), |
|
nn.Conv1d(channels, inter_channels, kernel_size=1, stride=1, padding=0), |
|
nn.BatchNorm1d(inter_channels), |
|
nn.ReLU(inplace=True), |
|
nn.Conv1d(inter_channels, channels, kernel_size=1, stride=1, padding=0), |
|
nn.BatchNorm1d(channels), |
|
) |
|
elif type == "2D": |
|
self.local_att = nn.Sequential( |
|
nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0), |
|
nn.BatchNorm2d(inter_channels), |
|
nn.ReLU(inplace=True), |
|
nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0), |
|
nn.BatchNorm2d(channels), |
|
) |
|
self.global_att = nn.Sequential( |
|
nn.AdaptiveAvgPool2d(1), |
|
nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0), |
|
nn.BatchNorm2d(inter_channels), |
|
nn.ReLU(inplace=True), |
|
nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0), |
|
nn.BatchNorm2d(channels), |
|
) |
|
else: |
|
raise f"the type is not supported." |
|
|
|
self.sigmoid = nn.Sigmoid() |
|
|
|
def forward(self, x, residual): |
|
flag = False |
|
xa = x + residual |
|
if xa.size(0) == 1: |
|
xa = torch.cat([xa, xa], dim=0) |
|
flag = True |
|
xl = self.local_att(xa) |
|
xg = self.global_att(xa) |
|
xlg = xl + xg |
|
wei = self.sigmoid(xlg) |
|
xo = 2 * x * wei + 2 * residual * (1 - wei) |
|
if flag: |
|
xo = xo[0].unsqueeze(0) |
|
return xo |
|
|
