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