|
|
import os |
|
|
from functools import partial |
|
|
import torch |
|
|
|
|
|
from .vmamba import VSSM |
|
|
from .csms6s import flops_selective_scan_fn,flops_selective_scan_ref |
|
|
|
|
|
|
|
|
def build_vssm_model(config, **kwargs): |
|
|
model_type = config.MODEL.TYPE |
|
|
if model_type in ["vssm"]: |
|
|
model = VSSM( |
|
|
patch_size=config.MODEL.VSSM.PATCH_SIZE, |
|
|
in_chans=config.MODEL.VSSM.IN_CHANS, |
|
|
num_classes=config.MODEL.NUM_CLASSES, |
|
|
depths=config.MODEL.VSSM.DEPTHS, |
|
|
dims=config.MODEL.VSSM.EMBED_DIM, |
|
|
|
|
|
ssm_d_state=config.MODEL.VSSM.SSM_D_STATE, |
|
|
ssm_ratio=config.MODEL.VSSM.SSM_RATIO, |
|
|
ssm_rank_ratio=config.MODEL.VSSM.SSM_RANK_RATIO, |
|
|
ssm_dt_rank=("auto" if config.MODEL.VSSM.SSM_DT_RANK == "auto" else int(config.MODEL.VSSM.SSM_DT_RANK)), |
|
|
ssm_act_layer=config.MODEL.VSSM.SSM_ACT_LAYER, |
|
|
ssm_conv=config.MODEL.VSSM.SSM_CONV, |
|
|
ssm_conv_bias=config.MODEL.VSSM.SSM_CONV_BIAS, |
|
|
ssm_drop_rate=config.MODEL.VSSM.SSM_DROP_RATE, |
|
|
ssm_init=config.MODEL.VSSM.SSM_INIT, |
|
|
forward_type=config.MODEL.VSSM.SSM_FORWARDTYPE, |
|
|
|
|
|
mlp_ratio=config.MODEL.VSSM.MLP_RATIO, |
|
|
mlp_act_layer=config.MODEL.VSSM.MLP_ACT_LAYER, |
|
|
mlp_drop_rate=config.MODEL.VSSM.MLP_DROP_RATE, |
|
|
|
|
|
drop_path_rate=config.MODEL.DROP_PATH_RATE, |
|
|
patch_norm=config.MODEL.VSSM.PATCH_NORM, |
|
|
norm_layer=config.MODEL.VSSM.NORM_LAYER, |
|
|
downsample_version=config.MODEL.VSSM.DOWNSAMPLE, |
|
|
patchembed_version=config.MODEL.VSSM.PATCHEMBED, |
|
|
gmlp=config.MODEL.VSSM.GMLP, |
|
|
use_checkpoint=config.TRAIN.USE_CHECKPOINT, |
|
|
|
|
|
posembed=config.MODEL.VSSM.POSEMBED, |
|
|
imgsize=config.DATA.IMG_SIZE, |
|
|
) |
|
|
return model |
|
|
|
|
|
return None |
|
|
|
|
|
|
|
|
def build_model(config, is_pretrain=False): |
|
|
model = None |
|
|
if model is None: |
|
|
model = build_vssm_model(config, is_pretrain) |
|
|
return model |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
