| | 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 |