GitHub - ReckonerInheritor/UFA-Fusion
UFA-FUSE: A novel deep supervised and hybrid model for multi-focus image fusion Yongsheng Zang, Dongming Zhou, Changcheng Wang, Rencan Nie, and Yanbu Gu
注意力机制
import torch
import torch.nn as nn
class Conv(nn.Module):
def __init__(self,out_ch):
super(Conv, self).__init__()
self.body=nn.Sequential(nn.Conv2d(out_ch, out_ch, 3, 1, 1),nn.LeakyReLU(negative_slope=1e-2,inplace=True))
def forward(self, x):
return self.body(x)
class SpatialAttention(nn.Module):
def __init__(self, kernel_size=7):
super(SpatialAttention, self).__init__()
assert kernel_size in (3, 7)
padding = 3 if kernel_size == 7 else 1
self.conv = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
self.sigmoid = nn.Sigmoid()
def forward(self, x):
avgout = torch.mean(x, dim=1, keepdim=True) #2,1,256,256
maxout, _= torch.max(x, dim=1, keepdim=True)#2,1,256,256
x = torch.cat([avgout, maxout], dim=1) #2,2,256,256
x = self.conv(x) #2,1,256,256
return x
class SpatialAttentionBlock(nn.Module):
def __init__(self):
super(SpatialAttentionBlock, self).__init__()
self.att1=SpatialAttention()
self.att2=SpatialAttention()
def forward(self,x1,x2):
EPSILON = 1e-10
att1=self.att1(x1) #2,1,256,256
att2=self.att2(x2)#2,1,256,256
mask1 = torch.exp(att1) / (torch.exp(att2) + torch.exp(att2) + EPSILON)#2,1,256,256
mask2 = torch.exp(att2) / (torch.exp(att1) + torch.exp(att2) + EPSILON)#2,1,256,256
x1_a=mask1*x1#2,64,256,256
x2_a=mask2*x2#2,64,256,256
return x1_a,x2_a
class ChannelAttention(nn.Module):
def __init__(self):
super(ChannelAttention, self).__init__()
self.avg=nn.AdaptiveAvgPool2d(1)
def forward(self,x):
return self.avg(x)
class ChannelAttentionBlock(nn.Module):
def __init__(self):
super(ChannelAttentionBlock, self).__init__()
self.ca1=nn.AdaptiveAvgPool2d(1)
self.ca2=nn.AdaptiveAvgPool2d(1)
def forward(self, x1,x2):
EPSILON = 1e-10
ca1=self.ca1(x1) #2,64,1,1
ca2 = self.ca1(x2) #2,64,1,1
mask1 = torch.exp(ca1) / (torch.exp(ca2) + torch.exp(ca1) + EPSILON)
mask2 = torch.exp(ca2) / (torch.exp(ca1) + torch.exp(ca2) + EPSILON)
x1_a=mask1*x1
x2_a=mask2*x2
return x1_a,x2_a
class attention(nn.Module):
def __init__(self,out_ch):
super(attention, self).__init__()
self.conv1=Conv(out_ch=out_ch)
self.conv2 = Conv(out_ch=out_ch)
self.att_ch=ChannelAttentionBlock()
self.att_sp=SpatialAttentionBlock()
def forward(self, x1,x2):
x1=self.conv1(x1) #2,64,256,256
x2=self.conv1(x2) #2,64,256,256
x1_c,x2_c=self.att_ch(x1,x2)
x1_s,x2_s=self.att_sp(x1_c,x2_c)#2,64,256,256
x1_out=self.conv2(x1_s)##2,64,256,256
x2_out=self.conv2(x2_s)##2,64,256,256
return x1_out,x2_out,x1,x2,x1_s,x2_s
class Resblock(nn.Module):
def __init__(self,out_ch):
super(Resblock, self).__init__()
self.body=nn.Sequential(nn.Conv2d(out_ch, out_ch, 3, 1, 1),nn.LeakyReLU(negative_slope=1e-2,inplace=True),
nn.Conv2d(out_ch, out_ch, 3, 1, 1))
def forward(self, x1,x2):
x1_out=self.body(x1)
x2_out=self.body(x2)
return x1_out,x2_out
class AF(nn.Module):
def __init__(self,args):
super(AF, self).__init__()
in_ch=args.in_ch
out_ch=args.out_ch
self.conv1 = nn.Sequential(nn.Conv2d(in_ch, out_ch, kernel_size=3,stride=1, padding=1, bias=True),nn.LeakyReLU(negative_slope=1e-2,inplace=True))
self.rb1=Resblock(out_ch)
self.rb2 = Resblock(out_ch)
self.rb3=Resblock(out_ch)
self.sc4 = attention(out_ch)
self.fusion=nn.Conv2d(out_ch*2,out_ch,1,1,0)
EX3 = [Conv(out_ch=out_ch) for i in range(4)]
self.EX3 = nn.Sequential(*EX3)
# conv
self.restruction= nn.Conv2d(out_ch, in_ch, kernel_size=3,stride=1, padding=1, bias=True)
self.sig = nn.Sigmoid()
def forward(self,x1,x2):
# x1,x2: 2,3,256,256
F1 =self.conv1(x1) #2,64,256,256
F2 =self.conv1(x2) #2,64,256,256
F1,F2=self.rb1(F1,F2) #2,64,256,256
F1,F2 = self.rb2(F1, F2) #2,64,256,256
F1,F2 = self.rb3(F1, F2)#2,64,256,256
F1,F2,F1_ori,F2_ori,F1_att,F2_att = self.sc4(F1, F2)
#F1,F2:2,64,256,256, #F1_ori,F2_ori 2,64,256,256 #F1_att,F2_att :2,64,256,256
fusion=self.fusion(torch.cat([F1,F2],dim=1))#2,64,256,256
fusion= self.EX3(fusion) #2,64,256,256
output = self.restruction(fusion)#2,3,256,256
output=self.sig(output) #2,3,256,256
return output,F1_ori,F2_ori,F1_att,F2_att
训练
def weights_init_xavier(m):
classname = m.__class__.__name__
if classname.find('Conv2d') != -1:
init.xavier_normal_(m.weight.data)
def training_setup_seed(seed):
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
np.random.seed(seed)
random.seed(seed)
torch.backends.cudnn.deterministic = True
def train(train_loader,args):
device=args.device
n_epochs=args.n_epochs
model_dir=args.saveModel_dir
batch_size=args.batch_size
train_num=15570
writer1 = SummaryWriter(log_dir="log/loss")
# ------- 1. define model --------
# define the net
net = AF(args).to(device)
net.apply(weights_init_xavier)
#define the loss
criterion = LpLssimLoss().to(args.device)
L1=torch.nn.L1Loss().to(args.device)
# MSE=torch.nn.MSELoss().to(args.device)
# ------- 2. define optimizer --------
optimizer = optim.Adam(net.parameters(), lr=0.0001)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=args.n_steps, gamma=args.gamma)
# ------- 3. training process --------
# ite_num = 0
running_loss = 0.0
# ite_num4val = 0
for epoch in range(0, n_epochs):
net.train()
t1 = time.time()
for i, (image1,image2,label) in enumerate(train_loader):
image1=image1.to(device)
image2 = image2.to(device)
label = label.to(device)
# ite_num = ite_num + 1
# ite_num4val = ite_num4val + 1
out,F1_ori,F2_ori,F1_a,F2_a=net(image1,image2)
l1=L1(out,label)
# mssim=1-ms_ssim( out, label, data_range=1, size_average=True)
ssim=criterion(out,label)
loss =0.8*ssim+0.2*l1
# if i==0:
# im1=image1[0].detach().cpu()
# im2=image2[0].detach().cpu()
# lab1=label[0].detach().cpu()
# out_image=out[0].detach().cpu()
# writer1.add_image('resutl', make_grid([im1,im2,lab1,out_image], nrow=5, padding=20, normalize=False, pad_value=1), epoch)
# writer1.add_image('F1_ori', make_grid(F1_ori[0].detach().cpu().unsqueeze(dim=1), nrow=8, padding=20, normalize=False, pad_value=1), epoch)
# writer1.add_image('F2_ori', make_grid(F2_ori[0].detach().cpu().unsqueeze(dim=1), nrow=8, padding=20, normalize=False, pad_value=1), epoch)
# writer1.add_image('F1_att', make_grid(F1_a[0].detach().cpu().unsqueeze(dim=1), nrow=8, padding=20, normalize=False, pad_value=1), epoch)
# writer1.add_image('F2_att', make_grid(F2_a[0].detach().cpu().unsqueeze(dim=1), nrow=8, padding=20, normalize=False, pad_value=1), epoch)
# writer1.add_image('fusion', make_grid(fusion[0].detach().cpu().unsqueeze(dim=1), nrow=8, padding=20, normalize=False, pad_value=1), epoch)
# writer1.add_image('F1_a', make_grid(F1_a[0].detach().cpu().unsqueeze(dim=1), nrow=8, padding=20, normalize=False, pad_value=1), epoch)
# writer1.add_image('F3_a', make_grid(F3_a[0].detach().cpu().unsqueeze(dim=1), nrow=8, padding=20, normalize=False, pad_value=1), epoch)
# writer1.add_image('atten_fusion', make_grid(atten_fusion[0].detach().cpu().unsqueeze(dim=1), nrow=8, padding=20, normalize=False, pad_value=1), epoch)
optimizer.zero_grad()
loss.backward()
optimizer.step()
# # print statistics
running_loss += loss.item()
print("[epoch: %3d/%3d, batch: %5d/%5d] train loss: %8f " % (epoch + 1, n_epochs, (i + 1) * batch_size, train_num, loss.item()))
writer1.add_scalar('训练损失', running_loss, epoch)
torch.save(net.state_dict(), model_dir + "epoch_%d_loss_%3f.pth" % (epoch, running_loss))
running_loss = 0.0
scheduler.step()
t2 = time.time()
print(t2-t1)
print('-------------Congratulations! Training Done!!!-------------')
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--in_ch", type=int, default=3,help='rgb is 3,gray is 1')
parser.add_argument("--out_ch", type=int, default=64)
parser.add_argument('--device', type=str, default='cuda:0')
parser.add_argument('--batch_size', type=int, default=2)
parser.add_argument('--n_epochs', type=int, default=6, help='number of epochs to train')
parser.add_argument('--n_steps', type=int, default=2, help='number of epochs to update learning rate')
parser.add_argument('--gamma', type=float, default=0.1, help='')
parser.add_argument('--dataset_dir', type=str, default="DataSet/TrainSet")
parser.add_argument('--saveModel_dir', type=str, default='experment_UFA/')
return parser.parse_args()
if __name__ == '__main__':
training_setup_seed(1)
args = parse_args()
transforms_ = [transforms.Resize((256,256), Image.BICUBIC),
# transforms.RandomHorizontalFlip(p=0.6),
transforms.ToTensor()]
train_set = SalObjDataset(dataset_dir=args.dataset_dir,transforms_=transforms_,rgb=True)
salobj_dataloader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True, num_workers=4)
train(salobj_dataloader, args)
CBAM的实现
CBAM将通道注意力机制和空间注意力机制进行一个结合，相比于SENet只关注通道的注意力机制可以取得更好的效果。其实现示意图如下所示，CBAM会对输入进来的特征层，分别进行通道注意力机制的处理和空间注意力机制的处理。
下图是通道注意力机制和空间注意力机制的具体实现方式：
图像的上半部分为通道注意力机制，通道注意力机制的实现可以分为两个部分，我们会对输入进来的单个特征层，分别进行全局平均池化和全局最大池化。之后对平均池化和最大池化的结果，利用共享的全连接层进行处理，我们会对处理后的两个结果进行相加，然后取一个sigmoid，此时我们获得了输入特征层每一个通道的权值（0-1之间）。在获得这个权值后，我们将这个权值乘上原输入特征层即可。
图像的下半部分为空间注意力机制，我们会对输入进来的特征层，在每一个特征点的通道上取最大值和平均值。之后将这两个结果进行一个堆叠，利用一次通道数为1的卷积调整通道数，然后取一个sigmoid，此时我们获得了输入特征层每一个特征点的权值（0-1之间）。在获得这个权值后，我们将这个权值乘上原输入特征层即可。
class ChannelAttention(nn.Module):
def __init__(self, in_planes, ratio=8):
super(ChannelAttention, self).__init__()
self.avg_pool = nn.AdaptiveAvgPool2d(1)
self.max_pool = nn.AdaptiveMaxPool2d(1)
# 利用1x1卷积代替全连接
self.fc1 = nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False)
self.relu1 = nn.ReLU()
self.fc2 = nn.Conv2d(in_planes // ratio, in_planes, 1, bias=False)
self.sigmoid = nn.Sigmoid()
def forward(self, x):
avg_out = self.fc2(self.relu1(self.fc1(self.avg_pool(x))))
max_out = self.fc2(self.relu1(self.fc1(self.max_pool(x))))
out = avg_out + max_out
return self.sigmoid(out)
class SpatialAttention(nn.Module):
def __init__(self, kernel_size=7):
super(SpatialAttention, self).__init__()
assert kernel_size in (3, 7), 'kernel size must be 3 or 7'
padding = 3 if kernel_size == 7 else 1
self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
self.sigmoid = nn.Sigmoid()
def forward(self, x):
avg_out = torch.mean(x, dim=1, keepdim=True)
max_out, _ = torch.max(x, dim=1, keepdim=True)
x = torch.cat([avg_out, max_out], dim=1)
x = self.conv1(x)
return self.sigmoid(x)
class cbam_block(nn.Module):
def __init__(self, channel, ratio=8, kernel_size=7):
super(cbam_block, self).__init__()
self.channelattention = ChannelAttention(channel, ratio=ratio)
self.spatialattention = SpatialAttention(kernel_size=kernel_size)
def forward(self, x):
x = x * self.channelattention(x)
x = x * self.spatialattention(x)
return x
SENet的实现
SENet是通道注意力机制的典型实现。
2017年提出的SENet是最后一届ImageNet竞赛的冠军，其实现示意图如下所示，对于输入进来的特征层，我们关注其每一个通道的权重，对于SENet而言，其重点是获得输入进来的特征层，每一个通道的权值。利用SENet，我们可以让网络关注它最需要关注的通道。
其具体实现方式就是：
1、对输入进来的特征层进行全局平均池化。
2、然后进行两次全连接，第一次全连接神经元个数较少，第二次全连接神经元个数和输入特征层相同。
3、在完成两次全连接后，我们再取一次Sigmoid将值固定到0-1之间，此时我们获得了输入特征层每一个通道的权值（0-1之间）。
4、在获得这个权值后，我们将这个权值乘上原输入特征层即可。
import torch
import torch.nn as nn
import math
class se_block(nn.Module):
def __init__(self, channel, ratio=16):
super(se_block, self).__init__()
self.avg_pool = nn.AdaptiveAvgPool2d(1)
self.fc = nn.Sequential(
nn.Linear(channel, channel // ratio, bias=False),
nn.ReLU(inplace=True),
nn.Linear(channel // ratio, channel, bias=False),
nn.Sigmoid()
)
def forward(self, x):
b, c, _, _ = x.size()
y = self.avg_pool(x).view(b, c)
y = self.fc(y).view(b, c, 1, 1)
return x * y
ECA的实现
ECANet是也是通道注意力机制的一种实现形式。ECANet可以看作是SENet的改进版。
ECANet的作者认为SENet对通道注意力机制的预测带来了副作用，捕获所有通道的依赖关系是低效并且是不必要的。
在ECANet的论文中，作者认为卷积具有良好的跨通道信息获取能力。
ECA模块的思想是非常简单的，它去除了原来SE模块中的全连接层，直接在全局平均池化之后的特征上通过一个1D卷积进行学习。
既然使用到了1D卷积，那么1D卷积的卷积核大小的选择就变得非常重要了，了解过卷积原理的同学很快就可以明白，1D卷积的卷积核大小会影响注意力机制每个权重的计算要考虑的通道数量。用更专业的名词就是跨通道交互的覆盖率。
如下图所示，左图是常规的SE模块，右图是ECA模块。ECA模块用1D卷积替换两次全连接。
————————————————
class eca_block(nn.Module):
def __init__(self, channel, b=1, gamma=2):
super(eca_block, self).__init__()
kernel_size = int(abs((math.log(channel, 2) + b) / gamma))
kernel_size = kernel_size if kernel_size % 2 else kernel_size + 1
self.avg_pool = nn.AdaptiveAvgPool2d(1)
self.conv = nn.Conv1d(1, 1, kernel_size=kernel_size, padding=(kernel_size - 1) // 2, bias=False)
self.sigmoid = nn.Sigmoid()
def forward(self, x):
y = self.avg_pool(x)
y = self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)
y = self.sigmoid(y)
return x * y.expand_as(x)
