396 lines
14 KiB
Python
396 lines
14 KiB
Python
import time
|
||
import math
|
||
import copy
|
||
from functools import partial
|
||
from typing import Optional, Callable
|
||
|
||
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
import torch.utils.checkpoint as checkpoint
|
||
from einops import rearrange, repeat
|
||
from timm.models.layers import DropPath, trunc_normal_
|
||
# from fvcore.nn import FlopCountAnalysis, flop_count_str, flop_count, parameter_count
|
||
|
||
DropPath.__repr__ = lambda self: f"timm.DropPath({self.drop_prob})"
|
||
|
||
# import mamba_ssm.selective_scan_fn (in which causal_conv1d is needed)
|
||
try:
|
||
from mamba_ssm.ops.selective_scan_interface import selective_scan_fn, selective_scan_ref
|
||
except:
|
||
pass
|
||
#
|
||
# # an alternative for mamba_ssm
|
||
# try:
|
||
# from selective_scan import selective_scan_fn as selective_scan_fn_v1
|
||
# from selective_scan import selective_scan_ref as selective_scan_ref_v1
|
||
# except:
|
||
# pass
|
||
|
||
"""
|
||
参考isdnet的主干，将其后三个阶段特征都不进行下采样
|
||
以2 3 306 306为例
|
||
|
||
"""
|
||
from mmseg.registry import MODELS
|
||
|
||
|
||
|
||
class PatchEmbed2D(nn.Module):
|
||
r""" Image to Patch Embedding
|
||
Args:
|
||
patch_size (int): Patch token size. Default: 4.
|
||
in_chans (int): Number of input image channels. Default: 3.
|
||
embed_dim (int): Number of linear projection output channels. Default: 96.
|
||
norm_layer (nn.Module, optional): Normalization layer. Default: None
|
||
"""
|
||
|
||
def __init__(self, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None, **kwargs):
|
||
super().__init__()
|
||
if isinstance(patch_size, int):
|
||
patch_size = (patch_size, patch_size)
|
||
self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, padding=1)
|
||
if norm_layer is not None:
|
||
self.norm = norm_layer(embed_dim)
|
||
else:
|
||
self.norm = None
|
||
|
||
def forward(self, x):
|
||
x = self.proj(x).permute(0, 2, 3, 1)
|
||
if self.norm is not None:
|
||
x = self.norm(x)
|
||
return x
|
||
|
||
|
||
|
||
class PatchMerging2D(nn.Module):
|
||
r""" Patch Merging Layer.
|
||
Args:
|
||
input_resolution (tuple[int]): Resolution of input feature.
|
||
dim (int): Number of input channels.
|
||
norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
|
||
"""
|
||
|
||
def __init__(self, dim, norm_layer=nn.LayerNorm):
|
||
super().__init__()
|
||
self.dim = dim
|
||
self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
|
||
self.norm = norm_layer(4 * dim)
|
||
|
||
# x 16 77 77 96
|
||
def forward(self, x):
|
||
#x = F.pad(x, pad=(0, 0, 1, 1, 1, 1), mode="constant", value=0) # 最后俩个填充第第一维度，中间两个填充第二维度，开头两个填充第一维度
|
||
B, H, W, C = x.shape
|
||
|
||
SHAPE_FIX = [-1, -1]
|
||
if (W % 2 != 0) or (H % 2 != 0):
|
||
#print(f"Warning, x.shape {x.shape} is not match even ===========", flush=True)
|
||
SHAPE_FIX[0] = H // 2
|
||
SHAPE_FIX[1] = W // 2
|
||
# 将空间信息融合进通道中
|
||
x0 = x[:, 0::2, 0::2, :] # B H/2 W/2 C
|
||
x1 = x[:, 1::2, 0::2, :] # B H/2 W/2 C
|
||
x2 = x[:, 0::2, 1::2, :] # B H/2 W/2 C
|
||
x3 = x[:, 1::2, 1::2, :] # B H/2 W/2 C
|
||
|
||
if SHAPE_FIX[0] > 0:
|
||
x0 = x0[:, :SHAPE_FIX[0], :SHAPE_FIX[1], :]
|
||
x1 = x1[:, :SHAPE_FIX[0], :SHAPE_FIX[1], :]
|
||
x2 = x2[:, :SHAPE_FIX[0], :SHAPE_FIX[1], :]
|
||
x3 = x3[:, :SHAPE_FIX[0], :SHAPE_FIX[1], :]
|
||
|
||
x = torch.cat([x0, x1, x2, x3], -1) # B H/2 W/2 4*C
|
||
x = x.view(B, H // 2, W // 2, 4 * C) # B H/2*W/2 4*C
|
||
|
||
x = self.norm(x) # 4c
|
||
x = self.reduction(x) # 2c
|
||
|
||
return x
|
||
|
||
|
||
|
||
from .mamba_block.mambair import mambairBlock
|
||
# from .mamba_block.vmamba import Vmamba_Block
|
||
from .mamba_block.vdb import VDB_Block
|
||
from .mamba_block.lfssblock import LFSSBlock
|
||
# from .mamba_block.my_efficent_block import VSSBlock
|
||
# from .mamba_block.my_efficent_block_td import VSSBlock
|
||
from .mamba_block.moganet import VSSBlock
|
||
class VSSLayer(nn.Module):
|
||
""" A basic Swin Transformer layer for one stage.
|
||
Args:
|
||
dim (int): Number of input channels.
|
||
depth (int): Number of blocks.
|
||
drop (float, optional): Dropout rate. Default: 0.0
|
||
attn_drop (float, optional): Attention dropout rate. Default: 0.0
|
||
drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
|
||
norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
|
||
downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
|
||
use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
dim,
|
||
depth,
|
||
attn_drop=0.,
|
||
drop_path=0.,
|
||
norm_layer=nn.LayerNorm,
|
||
downsample=None,
|
||
use_checkpoint=False,
|
||
d_state=16,
|
||
vssm_mode='',
|
||
**kwargs,
|
||
):
|
||
super().__init__()
|
||
self.dim = dim
|
||
self.use_checkpoint = use_checkpoint
|
||
self.vssm_mode=vssm_mode
|
||
|
||
self.blocks = nn.ModuleList([
|
||
mambairBlock(
|
||
hidden_dim=dim,
|
||
drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
|
||
norm_layer=norm_layer,
|
||
attn_drop_rate=attn_drop,
|
||
d_state=d_state,
|
||
vssm_mode= self.vssm_mode,
|
||
)
|
||
for i in range(depth)])
|
||
|
||
if True: # is this really applied? Yes, but been overriden later in VSSM!
|
||
def _init_weights(module: nn.Module):
|
||
for name, p in module.named_parameters():
|
||
if name in ["out_proj.weight"]:
|
||
p = p.clone().detach_() # fake init, just to keep the seed ....
|
||
nn.init.kaiming_uniform_(p, a=math.sqrt(5))
|
||
|
||
self.apply(_init_weights)
|
||
|
||
if downsample is not None:
|
||
self.downsample = downsample(dim=dim, norm_layer=norm_layer)
|
||
else:
|
||
self.downsample = None
|
||
|
||
def forward(self, x):
|
||
for blk in self.blocks:
|
||
if self.use_checkpoint:
|
||
x = checkpoint.checkpoint(blk, x)
|
||
else:
|
||
x = blk(x)
|
||
|
||
if self.downsample is not None:
|
||
x = self.downsample(x)
|
||
|
||
return x
|
||
|
||
|
||
|
||
|
||
|
||
@MODELS.register_module()
|
||
class VSSM(nn.Module):#原始ssm
|
||
def __init__(self, patch_size=4, in_chans=3, num_classes=4, depths=[2, 2, 9, 2],
|
||
dims=[96, 192, 384, 768], d_state=16, drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1,
|
||
norm_layer=nn.LayerNorm, patch_norm=True,vssm_mode='vssm',
|
||
use_checkpoint=False, final_upsample="expand_first",pretrain_model="", **kwargs):
|
||
super().__init__()
|
||
self.num_classes = num_classes
|
||
self.num_layers = len(depths)
|
||
if isinstance(dims, int):
|
||
dims = [int(dims * 2 ** i_layer) for i_layer in range(self.num_layers)]
|
||
self.embed_dim = dims[0]
|
||
self.num_features = dims[-1]
|
||
self.num_features_up = int(dims[0] * 2)
|
||
self.dims = dims
|
||
self.final_upsample = final_upsample
|
||
|
||
self.patch_embed = PatchEmbed2D(patch_size=patch_size, in_chans=in_chans, embed_dim=self.embed_dim,
|
||
norm_layer=norm_layer if patch_norm else None)
|
||
|
||
dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule
|
||
|
||
self.vssm_mode=vssm_mode
|
||
# build encoder and bottleneck layers
|
||
self.layers = nn.ModuleList()
|
||
for i_layer in range(self.num_layers): # 下采样层
|
||
layer = VSSLayer(
|
||
# dim=dims[i_layer], #int(embed_dim * 2 ** i_layer)
|
||
dim=int(dims[0] * 2 ** i_layer),
|
||
depth=depths[i_layer],
|
||
d_state=math.ceil(dims[0] / 6) if d_state is None else d_state, # 20240109
|
||
drop=drop_rate,
|
||
attn_drop=attn_drop_rate,
|
||
drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],
|
||
norm_layer=norm_layer,
|
||
downsample=PatchMerging2D if (i_layer < self.num_layers - 1) else None,
|
||
use_checkpoint=use_checkpoint,
|
||
vssm_mode=self.vssm_mode,
|
||
)
|
||
self.layers.append(layer)
|
||
|
||
# build decoder layers
|
||
|
||
self.norm = norm_layer(self.num_features)
|
||
|
||
|
||
if pretrain_model:
|
||
print('use pretrain model {}'.format(pretrain_model))
|
||
self.load_from(pretrain_model)
|
||
self.apply(self._init_weights)
|
||
else:
|
||
self.apply(self._init_weights)
|
||
|
||
def load_from(self, config):
|
||
# pretrained_path = config.MODEL.PRETRAIN_CKPT
|
||
pretrained_path = config
|
||
if pretrained_path is not None:
|
||
print("pretrained_path:{}".format(pretrained_path))
|
||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||
pretrained_dict = torch.load(pretrained_path, map_location=device)
|
||
if "model" not in pretrained_dict:
|
||
print("---start load pretrained modle by splitting---")
|
||
pretrained_dict = {k[17:]: v for k, v in pretrained_dict.items()}
|
||
for k in list(pretrained_dict.keys()):
|
||
if "output" in k:
|
||
print("delete key:{}".format(k))
|
||
del pretrained_dict[k]
|
||
msg = self.load_state_dict(pretrained_dict, strict=False)
|
||
# print(msg)
|
||
return
|
||
pretrained_dict = pretrained_dict['model']
|
||
print("---start load pretrained modle of swin encoder---")
|
||
|
||
model_dict = self.state_dict()
|
||
full_dict = copy.deepcopy(pretrained_dict)
|
||
for k, v in pretrained_dict.items():
|
||
if "layers." in k:
|
||
current_layer_num = 3 - int(k[7:8])
|
||
current_k = "layers_up." + str(current_layer_num) + k[8:]
|
||
full_dict.update({current_k: v})
|
||
for k in list(full_dict.keys()):
|
||
if k in model_dict:
|
||
if full_dict[k].shape != model_dict[k].shape:
|
||
print("delete:{};shape pretrain:{};shape model:{}".format(k, v.shape, model_dict[k].shape))
|
||
del full_dict[k]
|
||
|
||
msg = self.load_state_dict(full_dict, strict=False)
|
||
# print(msg)
|
||
else:
|
||
print("none pretrain")
|
||
|
||
|
||
def init_weight(self, pretrain_model):
|
||
|
||
state_dict = torch.load(pretrain_model)["state_dict"]#权重字典键
|
||
self_state_dict = self.state_dict()
|
||
for k, v in state_dict.items():
|
||
if k == 'features.0.conv.weight' and self.in_channels != 3:
|
||
v = torch.cat([v, v], dim=1)
|
||
self_state_dict.update({k: v})
|
||
self.load_state_dict(self_state_dict, strict=False)
|
||
def _init_weights(self, m: nn.Module):
|
||
"""
|
||
out_proj.weight which is previously initilized in VSSBlock, would be cleared in nn.Linear
|
||
no fc.weight found in the any of the model parameters
|
||
no nn.Embedding found in the any of the model parameters
|
||
so the thing is, VSSBlock initialization is useless
|
||
|
||
Conv2D is not intialized !!!
|
||
"""
|
||
# print(m, getattr(getattr(m, "weight", nn.Identity()), "INIT", None), isinstance(m, nn.Linear), "======================")
|
||
if isinstance(m, nn.Linear):
|
||
trunc_normal_(m.weight, std=.02)
|
||
if isinstance(m, nn.Linear) and m.bias is not None:
|
||
nn.init.constant_(m.bias, 0)
|
||
elif isinstance(m, nn.LayerNorm):
|
||
nn.init.constant_(m.bias, 0)
|
||
nn.init.constant_(m.weight, 1.0)
|
||
|
||
# Encoder and Bottleneck
|
||
def forward_features(self, x):
|
||
x = self.patch_embed(x) # 下采样4倍，通道数增加为96 ；24 56 56 96
|
||
|
||
x_downsample = []
|
||
for layer in self.layers:
|
||
x_downsample.append(x)
|
||
x = layer(x)
|
||
x = self.norm(x) # B H W C
|
||
|
||
x_middle=x_downsample[0:3]
|
||
x_middle.append(x)
|
||
return x, x_downsample,x_middle
|
||
|
||
# def forward_backbone(self, x):
|
||
# x = self.patch_embed(x)
|
||
|
||
# for layer in self.layers:
|
||
# x = layer(x)
|
||
# return x
|
||
|
||
# Dencoder and Skip connection
|
||
def forward_up_features(self, x, x_downsample):
|
||
for inx, layer_up in enumerate(self.layers_up):
|
||
if inx == 0:
|
||
x = layer_up(x)
|
||
else:
|
||
x = torch.cat([x, x_downsample[3 - inx]], -1)
|
||
x = self.concat_back_dim[inx](x)
|
||
x = layer_up(x)
|
||
|
||
x = self.norm_up(x) # B H W C
|
||
|
||
return x
|
||
|
||
def up_x4(self, x):
|
||
if self.final_upsample == "expand_first":
|
||
B, H, W, C = x.shape
|
||
x = self.up(x)
|
||
x = x.view(B, 4 * H, 4 * W, -1)
|
||
x = x.permute(0, 3, 1, 2) # B,C,H,W
|
||
x = self.output(x)
|
||
|
||
return x
|
||
|
||
def forward(self, x): # 24 3 224 224
|
||
x, x_downsample,x_middle= self.forward_features(x)
|
||
# x = self.forward_up_features(x,x_downsample)
|
||
# x = self.up_x4(x)
|
||
x_final=[]
|
||
for feat_every in x_middle:
|
||
feat=feat_every.permute(0,3,1,2)
|
||
x_final.append(feat)
|
||
return x_final
|
||
|
||
|
||
|
||
if __name__ == "__main__":
|
||
# check_vssm_equals_vmambadp()
|
||
# from torchstat import stat
|
||
#
|
||
# from thop import profile
|
||
|
||
weight="/media/cm/D450B76750B74ECC/Fusion_model/checkpoint/vmamba_tiny_e292.pth"
|
||
# model=MambaUnet_backbone(weight).to('cuda')
|
||
# # model = VSSM(use_checkpoint=weight).to('cuda')
|
||
# model.load_from(weight)
|
||
# int = torch.randn(2,1,224,224).cuda()
|
||
# out = model(int)
|
||
#
|
||
# # flops, params = profile(model, inputs=(int,))
|
||
# # print('flops:{}'.format(flops))
|
||
# # print('params:{}'.format(params))
|
||
# print(out.shape)
|
||
|
||
|
||
model = VSSM(pretrain_model=weight).to('cuda')
|
||
int = torch.randn(2, 3, 306, 306).cuda()
|
||
out = model(int)
|
||
|
||
# flops, params = profile(model, inputs=(int,))
|
||
# print('flops:{}'.format(flops))
|
||
# print('params:{}'.format(params))
|
||
|
||
# print(out.shape)
|
||
print(out)
|