| from __future__ import annotations
|
| import torch
|
| import torch.nn as nn
|
| import torch.nn.functional as F
|
| from monai.networks.layers.utils import get_act_layer
|
| import warnings
|
| warnings.filterwarnings("ignore")
|
| import math
|
| from functools import partial
|
| from typing import Callable
|
| from timm.models.layers import DropPath, to_2tuple
|
| from mamba_ssm.ops.selective_scan_interface import selective_scan_fn
|
| from einops import rearrange, repeat
|
|
|
|
|
| class CAB(nn.Module):
|
| def __init__(self, in_channels, out_channels=None, ratio=16, activation='relu'):
|
| super(CAB, self).__init__()
|
|
|
| self.in_channels = in_channels
|
| self.out_channels = out_channels
|
| if self.in_channels < ratio:
|
| ratio = self.in_channels
|
| self.reduced_channels = self.in_channels // ratio
|
| if self.out_channels == None:
|
| self.out_channels = in_channels
|
|
|
| self.avg_pool = nn.AdaptiveAvgPool2d(1)
|
| self.max_pool = nn.AdaptiveMaxPool2d(1)
|
| self.activation = get_act_layer(activation)
|
| self.fc1 = nn.Conv2d(self.in_channels, self.reduced_channels, 1, bias=False)
|
| self.fc2 = nn.Conv2d(self.reduced_channels, self.out_channels, 1, bias=False)
|
|
|
| self.sigmoid = nn.Sigmoid()
|
|
|
| nn.init.kaiming_normal_(self.fc1.weight, mode='fan_out', nonlinearity='relu')
|
| nn.init.kaiming_normal_(self.fc2.weight, mode='fan_out', nonlinearity='relu')
|
|
|
| def forward(self, x):
|
|
|
| avg = self.fc2(self.activation(self.fc1(self.avg_pool(x))))
|
| max = self.fc2(self.activation(self.fc1(self.max_pool(x))))
|
| attention = self.sigmoid(avg + max)
|
|
|
| return attention
|
|
|
| class SAB(nn.Module):
|
| def __init__(self, kernel_size=7):
|
| super(SAB, self).__init__()
|
| assert kernel_size in (3, 7, 11), "kernel_size must be 3, 7 or 11"
|
| padding = kernel_size // 2
|
|
|
| self.conv = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
|
| self.sigmoid = nn.Sigmoid()
|
| nn.init.kaiming_normal_(self.conv.weight, mode='fan_out', nonlinearity='relu')
|
|
|
| def forward(self, x):
|
| avg_out = torch.mean(x, dim=1, keepdim=True)
|
| max_out, _ = torch.max(x, dim=1, keepdim=True)
|
|
|
| x_cat = torch.cat([avg_out, max_out], dim=1)
|
| attention = self.sigmoid(self.conv(x_cat))
|
|
|
| return attention
|
|
|
|
|
|
|
|
|
| class ChannelAttention(nn.Module):
|
| """Channel attention used in RCAN.
|
| Args:
|
| num_feat (int): Channel number of intermediate features.
|
| squeeze_factor (int): Channel squeeze factor. Default: 16.
|
| """
|
|
|
| def __init__(self, num_feat, squeeze_factor=16):
|
| super(ChannelAttention, self).__init__()
|
| squeeze_channels = max(num_feat // squeeze_factor, 4)
|
| self.attention = nn.Sequential(
|
| nn.AdaptiveAvgPool2d(1),
|
| nn.Conv2d(num_feat, squeeze_channels, 1, padding=0),
|
| nn.ReLU(inplace=True),
|
| nn.Conv2d(squeeze_channels, num_feat, 1, padding=0),
|
| nn.Sigmoid())
|
|
|
| def forward(self, x):
|
| y = self.attention(x)
|
| return x * y
|
|
|
|
|
| class CAB(nn.Module):
|
| def __init__(self, num_feat, is_light_sr= False, compress_ratio=3,squeeze_factor=30):
|
| super(CAB, self).__init__()
|
| mid_channels = max(num_feat // compress_ratio, 4)
|
| if is_light_sr:
|
| self.cab = nn.Sequential(
|
| nn.Conv2d(num_feat, num_feat, 3, 1, 1, groups=num_feat),
|
| ChannelAttention(num_feat, squeeze_factor)
|
| )
|
| else:
|
| self.cab = nn.Sequential(
|
| nn.Conv2d(num_feat, mid_channels, 3, 1, 1),
|
| nn.GELU(),
|
| nn.Conv2d(mid_channels, num_feat, 3, 1, 1),
|
| ChannelAttention(num_feat, squeeze_factor)
|
| )
|
|
|
| def forward(self, x):
|
| return self.cab(x)
|
|
|
| class SS2D(nn.Module):
|
| def __init__(
|
| self,
|
| d_model,
|
| d_state=16,
|
| d_conv=3,
|
| expand=2.,
|
| dt_rank="auto",
|
| dt_min=0.001,
|
| dt_max=0.1,
|
| dt_init="random",
|
| dt_scale=1.0,
|
| dt_init_floor=1e-4,
|
| dropout=0.,
|
| conv_bias=True,
|
| bias=False,
|
| device=None,
|
| dtype=None,
|
| **kwargs,
|
| ):
|
| factory_kwargs = {"device": device, "dtype": dtype}
|
| super().__init__()
|
| self.d_model = d_model
|
| self.d_state = d_state
|
| self.d_conv = d_conv
|
| self.expand = expand
|
| self.d_inner = int(self.expand * self.d_model)
|
| self.dt_rank = math.ceil(self.d_model / 16) if dt_rank == "auto" else dt_rank
|
|
|
| self.in_proj = nn.Linear(self.d_model, self.d_inner * 2, bias=bias, **factory_kwargs)
|
| self.conv2d = nn.Conv2d(
|
| in_channels=self.d_inner,
|
| out_channels=self.d_inner,
|
| groups=self.d_inner,
|
| bias=conv_bias,
|
| kernel_size=d_conv,
|
| padding=(d_conv - 1) // 2,
|
| **factory_kwargs,
|
| )
|
| self.act = nn.SiLU()
|
|
|
| self.x_proj = (
|
| nn.Linear(self.d_inner, (self.dt_rank + self.d_state * 2), bias=False, **factory_kwargs),
|
| nn.Linear(self.d_inner, (self.dt_rank + self.d_state * 2), bias=False, **factory_kwargs),
|
| nn.Linear(self.d_inner, (self.dt_rank + self.d_state * 2), bias=False, **factory_kwargs),
|
| nn.Linear(self.d_inner, (self.dt_rank + self.d_state * 2), bias=False, **factory_kwargs),
|
| )
|
| self.x_proj_weight = nn.Parameter(torch.stack([t.weight for t in self.x_proj], dim=0))
|
| del self.x_proj
|
|
|
| self.dt_projs = (
|
| self.dt_init(self.dt_rank, self.d_inner, dt_scale, dt_init, dt_min, dt_max, dt_init_floor,
|
| **factory_kwargs),
|
| self.dt_init(self.dt_rank, self.d_inner, dt_scale, dt_init, dt_min, dt_max, dt_init_floor,
|
| **factory_kwargs),
|
| self.dt_init(self.dt_rank, self.d_inner, dt_scale, dt_init, dt_min, dt_max, dt_init_floor,
|
| **factory_kwargs),
|
| self.dt_init(self.dt_rank, self.d_inner, dt_scale, dt_init, dt_min, dt_max, dt_init_floor,
|
| **factory_kwargs),
|
| )
|
| self.dt_projs_weight = nn.Parameter(torch.stack([t.weight for t in self.dt_projs], dim=0))
|
| self.dt_projs_bias = nn.Parameter(torch.stack([t.bias for t in self.dt_projs], dim=0))
|
| del self.dt_projs
|
|
|
| self.A_logs = self.A_log_init(self.d_state, self.d_inner, copies=4, merge=True)
|
| self.Ds = self.D_init(self.d_inner, copies=4, merge=True)
|
|
|
| self.selective_scan = selective_scan_fn
|
|
|
| self.out_norm = nn.LayerNorm(self.d_inner)
|
| self.out_proj = nn.Linear(self.d_inner, self.d_model, bias=bias, **factory_kwargs)
|
| self.dropout = nn.Dropout(dropout) if dropout > 0. else None
|
|
|
| @staticmethod
|
| def dt_init(dt_rank, d_inner, dt_scale=1.0, dt_init="random", dt_min=0.001, dt_max=0.1, dt_init_floor=1e-4,
|
| **factory_kwargs):
|
| dt_proj = nn.Linear(dt_rank, d_inner, bias=True, **factory_kwargs)
|
|
|
|
|
| dt_init_std = dt_rank ** -0.5 * dt_scale
|
| if dt_init == "constant":
|
| nn.init.constant_(dt_proj.weight, dt_init_std)
|
| elif dt_init == "random":
|
| nn.init.uniform_(dt_proj.weight, -dt_init_std, dt_init_std)
|
| else:
|
| raise NotImplementedError
|
|
|
|
|
| dt = torch.exp(
|
| torch.rand(d_inner, **factory_kwargs) * (math.log(dt_max) - math.log(dt_min))
|
| + math.log(dt_min)
|
| ).clamp(min=dt_init_floor)
|
|
|
| inv_dt = dt + torch.log(-torch.expm1(-dt))
|
| with torch.no_grad():
|
| dt_proj.bias.copy_(inv_dt)
|
|
|
| dt_proj.bias._no_reinit = True
|
|
|
| return dt_proj
|
|
|
| @staticmethod
|
| def A_log_init(d_state, d_inner, copies=1, device=None, merge=True):
|
|
|
| A = repeat(
|
| torch.arange(1, d_state + 1, dtype=torch.float32, device=device),
|
| "n -> d n",
|
| d=d_inner,
|
| ).contiguous()
|
| A_log = torch.log(A)
|
| if copies > 1:
|
| A_log = repeat(A_log, "d n -> r d n", r=copies)
|
| if merge:
|
| A_log = A_log.flatten(0, 1)
|
| A_log = nn.Parameter(A_log)
|
| A_log._no_weight_decay = True
|
| return A_log
|
|
|
| @staticmethod
|
| def D_init(d_inner, copies=1, device=None, merge=True):
|
|
|
| D = torch.ones(d_inner, device=device)
|
| if copies > 1:
|
| D = repeat(D, "n1 -> r n1", r=copies)
|
| if merge:
|
| D = D.flatten(0, 1)
|
| D = nn.Parameter(D)
|
| D._no_weight_decay = True
|
| return D
|
|
|
| def forward_core(self, x: torch.Tensor):
|
| B, C, H, W = x.shape
|
| L = H * W
|
| K = 4
|
| x_hwwh = torch.stack([x.view(B, -1, L), torch.transpose(x, dim0=2, dim1=3).contiguous().view(B, -1, L)], dim=1).view(B, 2, -1, L)
|
| xs = torch.cat([x_hwwh, torch.flip(x_hwwh, dims=[-1])], dim=1)
|
|
|
| x_dbl = torch.einsum("b k d l, k c d -> b k c l", xs.view(B, K, -1, L), self.x_proj_weight)
|
| dts, Bs, Cs = torch.split(x_dbl, [self.dt_rank, self.d_state, self.d_state], dim=2)
|
| dts = torch.einsum("b k r l, k d r -> b k d l", dts.view(B, K, -1, L), self.dt_projs_weight)
|
| xs = xs.float().view(B, -1, L)
|
| dts = dts.contiguous().float().view(B, -1, L)
|
| Bs = Bs.float().view(B, K, -1, L)
|
| Cs = Cs.float().view(B, K, -1, L)
|
| Ds = self.Ds.float().view(-1)
|
| As = -torch.exp(self.A_logs.float()).view(-1, self.d_state)
|
| dt_projs_bias = self.dt_projs_bias.float().view(-1)
|
| out_y = self.selective_scan(
|
| xs, dts,
|
| As, Bs, Cs, Ds, z=None,
|
| delta_bias=dt_projs_bias,
|
| delta_softplus=True,
|
| return_last_state=False,
|
| ).view(B, K, -1, L)
|
| assert out_y.dtype == torch.float
|
|
|
| inv_y = torch.flip(out_y[:, 2:4], dims=[-1]).view(B, 2, -1, L)
|
| wh_y = torch.transpose(out_y[:, 1].view(B, -1, W, H), dim0=2, dim1=3).contiguous().view(B, -1, L)
|
| invwh_y = torch.transpose(inv_y[:, 1].view(B, -1, W, H), dim0=2, dim1=3).contiguous().view(B, -1, L)
|
|
|
| return out_y[:, 0], inv_y[:, 0], wh_y, invwh_y
|
|
|
| def forward(self, x: torch.Tensor, **kwargs):
|
| B, H, W, C = x.shape
|
|
|
| xz = self.in_proj(x)
|
| x, z = xz.chunk(2, dim=-1)
|
|
|
| x = x.permute(0, 3, 1, 2).contiguous()
|
| x = self.act(self.conv2d(x))
|
| y1, y2, y3, y4 = self.forward_core(x)
|
| assert y1.dtype == torch.float32
|
| y = y1 + y2 + y3 + y4
|
| y = torch.transpose(y, dim0=1, dim1=2).contiguous().view(B, H, W, -1)
|
| y = self.out_norm(y)
|
| y = y * F.silu(z)
|
| out = self.out_proj(y)
|
| if self.dropout is not None:
|
| out = self.dropout(out)
|
| return out
|
|
|
|
|
| class VSSBlock(nn.Module):
|
| def __init__(
|
| self,
|
| hidden_dim: int = 0,
|
| drop_path: float = 0,
|
| norm_layer: Callable[..., torch.nn.Module] = partial(nn.LayerNorm, eps=1e-6),
|
| attn_drop_rate: float = 0,
|
| d_state: int = 16,
|
| expand: float = 2.,
|
| is_light_sr: bool = False,
|
| **kwargs,
|
| ):
|
| super().__init__()
|
| self.ln_1 = norm_layer(hidden_dim)
|
| self.self_attention = SS2D(d_model=hidden_dim, d_state=d_state,expand=expand,dropout=attn_drop_rate, **kwargs)
|
| self.drop_path = DropPath(drop_path)
|
| self.skip_scale= nn.Parameter(torch.ones(hidden_dim))
|
| self.conv_blk = CAB(hidden_dim,is_light_sr)
|
| self.ln_2 = nn.LayerNorm(hidden_dim)
|
| self.skip_scale2 = nn.Parameter(torch.ones(hidden_dim))
|
|
|
|
|
|
|
| def forward(self, input, x_size):
|
|
|
| B, L, C = input.shape
|
| input = input.view(B, *x_size, C).contiguous()
|
| x = self.ln_1(input)
|
| x = input*self.skip_scale + self.drop_path(self.self_attention(x))
|
| x = x*self.skip_scale2 + self.conv_blk(self.ln_2(x).permute(0, 3, 1, 2).contiguous()).permute(0, 2, 3, 1).contiguous()
|
| x = x.view(B, -1, C).contiguous()
|
| return x
|
|
|
| class RoPE(nn.Module):
|
| def __init__(self, embed_dim, num_heads):
|
| super().__init__()
|
| self.head_dim = embed_dim // num_heads
|
| self.num_heads = num_heads
|
|
|
| def forward(self, x_size):
|
| H, W = x_size
|
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| pos_h = torch.arange(H, dtype=torch.float32, device=device)
|
| pos_w = torch.arange(W, dtype=torch.float32, device=device)
|
| inv_freq = 1.0 / (10000 ** (torch.arange(0, self.head_dim, 2, device=device).float() / self.head_dim))
|
| sin_h = torch.sin(torch.einsum("i,j->ij", pos_h, inv_freq))
|
| cos_h = torch.cos(torch.einsum("i,j->ij", pos_h, inv_freq))
|
| sin_w = torch.sin(torch.einsum("i,j->ij", pos_w, inv_freq))
|
| cos_w = torch.cos(torch.einsum("i,j->ij", pos_w, inv_freq))
|
| sin = torch.einsum("i,j->ij", sin_h[:, 0], sin_w[:, 0]).unsqueeze(0).unsqueeze(0)
|
| cos = torch.einsum("i,j->ij", cos_h[:, 0], cos_w[:, 0]).unsqueeze(0).unsqueeze(0)
|
| sin = sin.expand(self.num_heads, -1, -1, -1).contiguous()
|
| cos = cos.expand(self.num_heads, -1, -1, -1).contiguous()
|
| return sin, cos
|
|
|
| def rotate_every_two(x):
|
| if x.shape[-1] % 2 != 0:
|
| x = F.pad(x, (0, 1), mode='constant', value=0)
|
| pad = True
|
| else:
|
| pad = False
|
|
|
| x1 = x[..., ::2]
|
| x2 = x[..., 1::2]
|
| out = torch.stack((-x2, x1), -1).reshape(*x.shape[:-1], -1)
|
|
|
| return out[..., :x.shape[-1]-1] if pad else out
|
|
|
|
|
| def theta_shift(x, sin, cos):
|
| if sin.shape[-1] < x.shape[-1]:
|
| pad = x.shape[-1] - sin.shape[-1]
|
| sin = F.pad(sin, (0, pad), mode='constant', value=0)
|
| cos = F.pad(cos, (0, pad), mode='constant', value=1)
|
| elif sin.shape[-1] > x.shape[-1]:
|
| sin = sin[..., :x.shape[-1]]
|
| cos = cos[..., :x.shape[-1]]
|
|
|
| return (x * cos) + (rotate_every_two(x) * sin)
|
|
|
| class OverlapWindowAttention(nn.Module):
|
| def __init__(self, dim, num_heads=4, window_size=7, shift_size=3):
|
| super().__init__()
|
| self.dim = dim
|
| self.num_heads = num_heads
|
| self.head_dim = dim // num_heads
|
| self.scale = self.head_dim ** -0.5
|
| self.window_size = window_size
|
| self.shift_size = shift_size
|
| self.qkv = nn.Conv2d(dim, dim * 3, kernel_size=1)
|
| self.proj = nn.Conv2d(dim, dim, kernel_size=1)
|
|
|
| def forward(self, x, sin, cos):
|
|
|
| B, C, H, W = x.shape
|
| ws = self.window_size
|
| pad_h = (ws - H % ws) % ws
|
| pad_w = (ws - W % ws) % ws
|
| x = F.pad(x, (0, pad_w, 0, pad_h), mode='reflect')
|
| H_pad, W_pad = x.shape[2], x.shape[3]
|
|
|
| if self.shift_size > 0:
|
| x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(2, 3))
|
|
|
| qkv = self.qkv(x)
|
| qkv = rearrange(qkv, 'b (m c) h w -> m b c h w', m=3)
|
| q, k, v = qkv[0], qkv[1], qkv[2]
|
| q = q.view(B, self.num_heads, self.head_dim, H_pad, W_pad)
|
| k = k.view(B, self.num_heads, self.head_dim, H_pad, W_pad)
|
| v = v.view(B, self.num_heads, self.head_dim, H_pad, W_pad)
|
| q = theta_shift(q, sin, cos) * self.scale
|
| k = theta_shift(k, sin, cos)
|
|
|
| q = q.view(B, C, H_pad, W_pad)
|
| k = k.view(B, C, H_pad, W_pad)
|
| v = v.view(B, C, H_pad, W_pad)
|
|
|
| q = rearrange(q, 'b c (h ws1) (w ws2) -> b (h w) (ws1 ws2) c', ws1=ws, ws2=ws)
|
| k = rearrange(k, 'b c (h ws1) (w ws2) -> b (h w) (ws1 ws2) c', ws1=ws, ws2=ws)
|
| v = rearrange(v, 'b c (h ws1) (w ws2) -> b (h w) (ws1 ws2) c', ws1=ws, ws2=ws)
|
|
|
| B, num_windows, window_len, C_new = q.shape
|
| assert C_new % self.num_heads == 0, f"C_new={C_new} 不能整除 num_heads={self.num_heads}"
|
| head_dim_new = C_new // self.num_heads
|
|
|
| q = q.view(B, num_windows, window_len, self.num_heads, head_dim_new).transpose(2, 3)
|
| k = k.view(B, num_windows, window_len, self.num_heads, head_dim_new).transpose(2, 3)
|
| v = v.view(B, num_windows, window_len, self.num_heads, head_dim_new).transpose(2, 3)
|
|
|
| attn = torch.softmax(q @ k.transpose(-2, -1), dim=-1)
|
| out = (attn @ v).transpose(2, 3).reshape(B, num_windows, window_len, self.num_heads * head_dim_new)
|
| out = rearrange(out, 'b (h w) (ws1 ws2) c -> b c (h ws1) (w ws2)', h=H_pad // ws, ws1=ws, ws2=ws, w=W_pad // ws)
|
|
|
| if self.shift_size > 0:
|
| out = torch.roll(out, shifts=(self.shift_size, self.shift_size), dims=(2, 3))
|
|
|
| out = out[:, :, :H, :W]
|
| out = self.proj(out)
|
| return out
|
|
|
| class ShallowFusionAttnBlock(nn.Module):
|
| def __init__(self, dim, num_heads=4, window_size=7, shift_size=3):
|
| super().__init__()
|
| self.dim = dim
|
| self.attn = OverlapWindowAttention(dim, num_heads=num_heads, window_size=window_size, shift_size=shift_size)
|
| self.rope = RoPE(embed_dim=dim, num_heads=num_heads)
|
| self.conv1 = nn.Conv2d(dim * 2, dim, kernel_size=3, padding=1)
|
| self.conv2 = nn.Conv2d(dim * 2, dim, kernel_size=3, padding=1)
|
| self.vss = VSSBlock(dim)
|
|
|
| def patch_unembed(self, x, h, w):
|
| return x.transpose(1, 2).reshape(x.size(0), -1, h, w)
|
|
|
| def patch_embed(self, x):
|
| return x.flatten(2).transpose(1, 2)
|
|
|
| def forward(self, I1, I2, h, w):
|
| B, C, H, W = I1.shape
|
|
|
| diff = torch.abs(I1 - I2)
|
| H_pad = (self.attn.window_size - h % self.attn.window_size) % self.attn.window_size + h
|
| W_pad = (self.attn.window_size - w % self.attn.window_size) % self.attn.window_size + w
|
| sin, cos = self.rope((H_pad, W_pad))
|
| diff_attn = self.attn(diff, sin, cos)
|
| token_attn = self.patch_embed(diff_attn)
|
| I1_token = self.patch_embed(I1)
|
| I2_token = self.patch_embed(I2)
|
| I1 = I1_token + token_attn
|
| I2 = I2_token + token_attn
|
|
|
| I1_un = self.patch_unembed(I1, h, w)
|
| I2_un = self.patch_unembed(I2, h, w)
|
|
|
| I1_local = self.conv1(torch.cat([I1_un, I2_un], dim=1)) + I1_un
|
| I2_local = self.conv2(torch.cat([I2_un, I1_un], dim=1)) + I2_un
|
|
|
| I1_token = self.patch_embed(I1_local)
|
| I2_token = self.patch_embed(I2_local)
|
| vss_feat_1 = self.vss(I1_token, (h, w)).transpose(1, 2).view(B, C, h, w)
|
| vss_feat_2 = self.vss(I2_token, (h, w)).transpose(1, 2).view(B, C, h, w)
|
|
|
| I1_fuse = I1_local + vss_feat_1
|
| I2_fuse = I2_local + vss_feat_2
|
| return I1_fuse, I2_fuse |
