41 lines
1.3 KiB
Python
41 lines
1.3 KiB
Python
"""Basic building blocks: SABlock (Spatial Attention), SEBlock (Squeeze-and-Excitation)
|
|
|
|
Original source: D-DFFNet/model/ours/blocks.py
|
|
"""
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
class SABlock(nn.Module):
|
|
"""Spatial self-attention block"""
|
|
def __init__(self, in_channels, out_channels):
|
|
super(SABlock, self).__init__()
|
|
self.attention = nn.Sequential(
|
|
nn.Conv2d(in_channels, out_channels, 3, padding=1, bias=False),
|
|
nn.Sigmoid()
|
|
)
|
|
self.conv = nn.Conv2d(in_channels, out_channels, 3, padding=1, bias=False)
|
|
|
|
def forward(self, x):
|
|
attention_mask = self.attention(x)
|
|
features = self.conv(x)
|
|
return torch.mul(features, attention_mask)
|
|
|
|
|
|
class SEBlock(nn.Module):
|
|
"""Squeeze-and-excitation block"""
|
|
def __init__(self, channels, r=16):
|
|
super(SEBlock, self).__init__()
|
|
self.r = r
|
|
self.squeeze = nn.Sequential(
|
|
nn.Linear(channels, channels // self.r),
|
|
nn.ReLU(),
|
|
nn.Linear(channels // self.r, channels),
|
|
nn.Sigmoid()
|
|
)
|
|
|
|
def forward(self, x):
|
|
B, C, H, W = x.size()
|
|
squeeze = self.squeeze(torch.mean(x, dim=(2, 3))).view(B, C, 1, 1)
|
|
return torch.mul(x, squeeze)
|