207 lines
7.8 KiB
Python
207 lines
7.8 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from timm.models.layers import DropPath
|
|
from .models.modules import MixedAttentionBlock
|
|
|
|
|
|
class decoder(nn.Module):
|
|
r""" Multistage decoder.
|
|
|
|
Args:
|
|
embed_dim (int): Dimension for attention. Default 384
|
|
dim (int): Patch embedding dimension. Default 96
|
|
img_size (int): Input image size. Default 224
|
|
mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
|
|
"""
|
|
|
|
def __init__(self, embed_dim=384, dims=[96, 192, 384], img_size=224, mlp_ratio=3):
|
|
super(decoder, self).__init__()
|
|
self.img_size = img_size
|
|
self.dims = dims
|
|
self.embed_dim = embed_dim
|
|
self.fusion1 = multiscale_fusion(in_dim=dims[2], f_dim=dims[1], kernel_size=(3, 3),
|
|
img_size=(img_size // 8, img_size // 8), stride=(2, 2), padding=(1, 1))
|
|
self.fusion2 = multiscale_fusion(in_dim=dims[1], f_dim=dims[0], kernel_size=(3, 3),
|
|
img_size=(img_size // 4, img_size // 4), stride=(2, 2), padding=(1, 1))
|
|
self.fusion3 = multiscale_fusion(in_dim=dims[0], f_dim=dims[0], kernel_size=(7, 7),
|
|
img_size=(img_size // 1, img_size // 1), stride=(4, 4), padding=(2, 2),
|
|
fuse=False)
|
|
|
|
self.mixatt1 = MixedAttention(in_dim=dims[1], dim=embed_dim, img_size=(img_size // 8, img_size // 8),
|
|
window_size=(img_size // 32), num_heads=1, mlp_ratio=mlp_ratio, depth=2)
|
|
self.mixatt2 = MixedAttention(in_dim=dims[0], dim=embed_dim, img_size=(img_size // 4, img_size // 4),
|
|
window_size=(img_size // 32), num_heads=1, mlp_ratio=mlp_ratio, depth=2)
|
|
|
|
self.proj1 = nn.Linear(dims[2], 1)
|
|
self.proj2 = nn.Linear(dims[1], 1)
|
|
self.proj3 = nn.Linear(dims[0], 1)
|
|
self.proj4 = nn.Linear(dims[0], 1)
|
|
|
|
def forward(self, f):
|
|
fea_1_16, fea_1_8, fea_1_4 = f # fea_1_16:1/16
|
|
B, _, _ = fea_1_16.shape
|
|
fea_1_8 = self.fusion1(fea_1_16, fea_1_8)
|
|
fea_1_8 = self.mixatt1(fea_1_8)
|
|
fea_1_4 = self.fusion2(fea_1_8, fea_1_4)
|
|
fea_1_4 = self.mixatt2(fea_1_4)
|
|
fea_1_1 = self.fusion3(fea_1_4)
|
|
fea_1_16 = self.proj1(fea_1_16)
|
|
mask_1_16 = fea_1_16.transpose(1, 2).reshape(B, 1, self.img_size // 16, self.img_size // 16)
|
|
fea_1_8 = self.proj2(fea_1_8)
|
|
mask_1_8 = fea_1_8.transpose(1, 2).reshape(B, 1, self.img_size // 8, self.img_size // 8)
|
|
fea_1_4 = self.proj3(fea_1_4)
|
|
mask_1_4 = fea_1_4.transpose(1, 2).reshape(B, 1, self.img_size // 4, self.img_size // 4)
|
|
fea_1_1 = self.proj4(fea_1_1)
|
|
mask_1_1 = fea_1_1.transpose(1, 2).reshape(B, 1, self.img_size // 1, self.img_size // 1)
|
|
return [mask_1_16, mask_1_8, mask_1_4, mask_1_1]
|
|
|
|
def flops(self):
|
|
flops = 0
|
|
flops += self.fusion1.flops()
|
|
flops += self.fusion2.flops()
|
|
flops += self.fusion3.flops()
|
|
flops += self.mixatt1.flops()
|
|
flops += self.mixatt2.flops()
|
|
|
|
flops += self.img_size // 16 * self.img_size // 16 * self.dims[2]
|
|
flops += self.img_size // 8 * self.img_size // 8 * self.dims[1]
|
|
flops += self.img_size // 4 * self.img_size // 4 * self.dims[0]
|
|
flops += self.img_size // 1 * self.img_size // 1 * self.dims[0]
|
|
|
|
return flops
|
|
|
|
|
|
class multiscale_fusion(nn.Module):
|
|
r""" Upsampling and feature fusion.
|
|
|
|
Args:
|
|
in_dim (int): Number of input feature channels.
|
|
f_dim (int): Number of fusion feature channels.
|
|
img_size (int): Image size after upsampling.
|
|
kernel_size (tuple(int)): The size of the sliding blocks.
|
|
stride (int): The stride of the sliding blocks in the input spatial dimensions, can be regarded as upsampling ratio.
|
|
padding (int): Implicit zero padding to be added on both sides of input.
|
|
fuse (bool): If True, concat features from different levels.
|
|
"""
|
|
|
|
def __init__(self, in_dim, f_dim, kernel_size, img_size, stride, padding, fuse=True):
|
|
super(multiscale_fusion, self).__init__()
|
|
self.fuse = fuse
|
|
self.norm = nn.LayerNorm(in_dim)
|
|
self.in_dim = in_dim
|
|
self.f_dim = f_dim
|
|
self.kernel_size = kernel_size
|
|
self.img_size = img_size
|
|
self.project = nn.Linear(in_dim, in_dim * kernel_size[0] * kernel_size[1])
|
|
self.upsample = nn.Fold(output_size=img_size, kernel_size=kernel_size, stride=stride, padding=padding)
|
|
if self.fuse:
|
|
self.mlp1 = nn.Sequential(
|
|
nn.Linear(in_dim + f_dim, f_dim),
|
|
nn.GELU(),
|
|
nn.Linear(f_dim, f_dim),
|
|
)
|
|
else:
|
|
self.proj = nn.Linear(in_dim, f_dim)
|
|
|
|
def forward(self, fea, fea_1=None):
|
|
fea = self.project(self.norm(fea))
|
|
fea = self.upsample(fea.transpose(1, 2))
|
|
B, C, _, _ = fea.shape
|
|
fea = fea.view(B, C, -1).transpose(1, 2) # .contiguous()
|
|
if self.fuse:
|
|
fea = torch.cat([fea, fea_1], dim=2)
|
|
fea = self.mlp1(fea)
|
|
else:
|
|
fea = self.proj(fea)
|
|
return fea
|
|
|
|
def flops(self):
|
|
N = self.img_size[0] * self.img_size[1]
|
|
flops = 0
|
|
# norm
|
|
flops += N * self.in_dim
|
|
# proj
|
|
flops += N * self.in_dim * self.in_dim * self.kernel_size[0] * self.kernel_size[1]
|
|
# mlp
|
|
flops += N * (self.in_dim + self.f_dim) * self.f_dim
|
|
flops += N * self.f_dim * self.f_dim
|
|
return flops
|
|
|
|
|
|
class MixedAttention(nn.Module):
|
|
r""" Mixed Attention Module.
|
|
|
|
Args:
|
|
in_dim (int): Number of input feature channels.
|
|
dim (int): Number for attention.
|
|
img_size (int): Image size after upsampling.
|
|
num_heads (int): Number of attention heads.
|
|
mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
|
|
depth (int): The number of MAB stacked.
|
|
"""
|
|
|
|
def __init__(self, in_dim, dim, img_size, window_size, num_heads=1, mlp_ratio=4, depth=2, drop_path=0.):
|
|
super(MixedAttention, self).__init__()
|
|
|
|
self.img_size = img_size
|
|
self.in_dim = in_dim
|
|
self.dim = dim
|
|
self.norm1 = nn.LayerNorm(in_dim)
|
|
self.mlp1 = nn.Sequential(
|
|
nn.Linear(in_dim, dim),
|
|
nn.GELU(),
|
|
nn.Linear(dim, dim),
|
|
)
|
|
self.blocks = nn.ModuleList([
|
|
MixedAttentionBlock(dim=dim, img_size=img_size, window_size=window_size, num_heads=num_heads,
|
|
mlp_ratio=mlp_ratio)
|
|
for i in range(depth)])
|
|
self.norm2 = nn.LayerNorm(dim)
|
|
self.mlp2 = nn.Sequential(
|
|
nn.Linear(dim, in_dim),
|
|
nn.GELU(),
|
|
nn.Linear(in_dim, in_dim),
|
|
)
|
|
self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
|
|
|
|
def forward(self, fea):
|
|
fea = self.mlp1(self.norm1(fea))
|
|
for blk in self.blocks:
|
|
fea = blk(fea)
|
|
fea = self.drop_path(self.mlp2(self.norm2(fea)))
|
|
return fea
|
|
|
|
def flops(self):
|
|
flops = 0
|
|
N = self.img_size[0] * self.img_size[1]
|
|
# norm1
|
|
flops += N * self.in_dim
|
|
# mlp1
|
|
flops += N * self.in_dim * self.dim
|
|
flops += N * self.dim * self.dim
|
|
# blks
|
|
for blk in self.blocks:
|
|
flops += blk.flops()
|
|
# norm2
|
|
flops += N * self.dim
|
|
# mlp2
|
|
flops += N * self.in_dim * self.dim
|
|
flops += N * self.dim * self.dim
|
|
return flops
|
|
|
|
|
|
if __name__ == '__main__':
|
|
# Test
|
|
model = decoder(embed_dim=384, dim=96, img_size=224)
|
|
model.cuda()
|
|
f = []
|
|
f.append(torch.randn((1, 196, 384)).cuda())
|
|
f.append(torch.randn((1, 784, 192)).cuda())
|
|
f.append(torch.randn((1, 3136, 96)).cuda())
|
|
|
|
y = model(f)
|
|
print(y[0].shape)
|
|
print(y[1].shape)
|
|
print(y[2].shape)
|
|
print(y[3].shape)
|