在# Copyright © 2015-present, Facebook, Inc.
# All rights reserved.
####################################### 等同于 vit_ce #############################
import torch
import torch.nn as nn
from functools import partial
from torch import Tensor
from typing import Optional

from timm.models.vision_transformer import VisionTransformer, cfg
from timm.models.registry import register_model
from timm.models.layers import trunc_normal, lecun_normal_

from timm.models.layers import DropPath, to_2tuple
from timm.models.vision_transformer import load_weights

import math

from collections import namedtuple

from mamba_ssm.modules.mamba_simple import Mamba
from mamba_ssm.utils.generation import GenerationMixin
from mamba_ssm.utils.hf import load_config_hf, load_state_dict_hf
#################### 结合当前项目引入 #########################
from lib.models.mamba_fetrack.rope import
import random
from lib.models.layers.head import build_box_head
import importlib
import lib.train.admin.settings as ws_settings
from .utils import combine_tokens, recover_tokens
#################################################################
try:
from mamba_ssm.ops.triton.layernorm import RMSNorm, layer_norm_fn, rms_norm_fn
except ImportError:
RMSNorm, layer_norm_fn, rms_norm_fn = None, None, None


all = [
‘vim_tiny_patch16_224’, ‘vim_small_patch16_224’, ‘vim_base_patch16_224’,
‘vim_tiny_patch16_384’, ‘vim_small_patch16_384’, ‘vim_base_patch16_384’,
]

# 原始数据 B,C,H,W → 切分成小方块 patches → Batch_size, embedding_dim, num_patches
class PatchEmbed(nn.Module):
“”" 2D Image to Patch Embedding
“”"
def init(self, img_size=224, patch_size=16, stride=16, in_chans=3, embed_dim=768, norm_layer=None, flatten=True):
super().init()
img_size = to_2tuple(img_size)
patch_size = to_2tuple(patch_size)
self.img_size = img_size
self.patch_size = patch_size
self.grid_size = ((img_size[0] - patch_size[0]) // stride + 1, (img_size[1] - patch_size[1]) // stride + 1)
self.num_patches = self.grid_size[0] * self.grid_size[1]
self.flatten = flatten
# proj就是切块操作
self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride)
self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()

def forward(self, x):
# print(‘x.shape:’,x.shape)
B, C, H, W = x.shape
# assert H == self.img_size[0] and W == self.img_size[1], <br/> # f"Input image size ({H}{W}) doesn’t match model ({self.img_size[0]}*{self.img_size[1]}).“
x = self.proj(x) # 将图像切分为小方块,尺寸就是patch_size
# 切块之后,维度： batch_size, embedding_dim, grid_size, grid_size
# 原始 （B,C,H,W）–> (B,N,grid_size,grid_size)
if self.flatten: # 将小方块拉平,构成序列
# 展平操作：batch_size, embedding_dim, num_patches
x = x.flatten(2).transpose(1, 2) # BCHW -> BNC
x = self.norm(x)
return x

# 定义mamba的encoder模块
class Block(nn.Module):
def init(
self, dim, mixer_cls, norm_cls=nn.LayerNorm, fused_add_norm=False, residual_in_fp32=False,drop_path=0.,
):
“””
Simple block wrapping a mixer class with LayerNorm/RMSNorm and residual connection"

This Block has a slightly different structure compared to a regular
prenorm Transformer block.
The standard block is: LN -> MHA/MLP -> Add.
[Ref: https://arxiv.org/abs/2002.04745]
Here we have: Add -> LN -> Mixer, returning both
the hidden_states (output of the mixer) and the residual.
This is purely for performance reasons, as we can fuse add and LayerNorm.
The residual needs to be provided (except for the very first block).
“”“
super().init()
self.residual_in_fp32 = residual_in_fp32
self.fused_add_norm = fused_add_norm
self.mixer = mixer_cls(dim) # 混合方式
self.norm = norm_cls(dim)
self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
if self.fused_add_norm:
assert RMSNorm is not None, “RMSNorm import fails”
assert isinstance(
self.norm, (nn.LayerNorm, RMSNorm)
), “Only LayerNorm and RMSNorm are supported for fused_add_norm”

def forward(
self, hidden_states: Tensor, residual: Optional[Tensor] = None, inference_params=None
):
r”"“Pass the input through the encoder layer.

Args:
hidden_states: the sequence to the encoder layer (required).
residual: hidden_states = Mixer(LN(residual))
“””
if not self.fused_add_norm:
if residual is None:
residual = hidden_states
else:
residual = residual + self.drop_path(hidden_states)

hidden_states = self.norm(residual.to(dtype=self.norm.weight.dtype))
if self.residual_in_fp32:
residual = residual.to(torch.float32)
else:
fused_add_norm_fn = rms_norm_fn if isinstance(self.norm, RMSNorm) else layer_norm_fn
if residual is None:
hidden_states, residual = fused_add_norm_fn(
hidden_states,
self.norm.weight,
self.norm.bias,
residual=residual,
prenorm=True,
residual_in_fp32=self.residual_in_fp32,
eps=self.norm.eps,
)
else:
hidden_states, residual = fused_add_norm_fn(
self.drop_path(hidden_states), # 随机失活之后保证模型具有更高的鲁棒性
self.norm.weight,
self.norm.bias,
residual=residual,
prenorm=True,
residual_in_fp32=self.residual_in_fp32,
eps=self.norm.eps,
)
hidden_states = self.mixer(hidden_states, inference_params=inference_params)
return hidden_states, residual

def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
return self.mixer.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype, **kwargs)


def create_block(
d_model,
ssm_cfg=None,
norm_epsilon=1e-5,
drop_path=0.,
rms_norm=False,
residual_in_fp32=False,
fused_add_norm=False,
layer_idx=None,
device=None,
dtype=None,
if_bimamba=False, # 是否使用bimamba,使用时就会采用 forward 和 backward
bimamba_type=“none”,
if_devide_out=False,
init_layer_scale=None,
):
if if_bimamba:
bimamba_type = “v1”
if ssm_cfg is None:
ssm_cfg = {}
factory_kwargs = {“device”: device, “dtype”: dtype}
# cls_tokens可以拼接到不同位置上
mixer_cls = partial(Mamba, layer_idx=layer_idx, bimamba_type=bimamba_type, if_devide_out=if_devide_out, init_layer_scale=init_layer_scale, **ssm_cfg, **factory_kwargs)
norm_cls = partial(
nn.LayerNorm if not rms_norm else RMSNorm, eps=norm_epsilon, **factory_kwargs
)
block = Block(
d_model,
mixer_cls,
norm_cls=norm_cls,
drop_path=drop_path,
fused_add_norm=fused_add_norm,
residual_in_fp32=residual_in_fp32,
)
block.layer_idx = layer_idx
return block


# https://github.com/huggingface/transformers/blob/c28d04e9e252a1a099944e325685f14d242ecdcd/src/transformers/models/gpt2/modeling_gpt2.py#L454
def init_weights(
module,
n_layer,
initializer_range=0.02, # Now only used for embedding layer.
rescale_prenorm_residual=True,
n_residuals_per_layer=1, # Change to 2 if we have MLP
):
if isinstance(module, nn.Linear):
if module.bias is not None:
if not getattr(module.bias, "no_reinit", False):
nn.init.zeros(module.bias)
elif isinstance(module, nn.Embedding):
nn.init.normal(module.weight, std=initializer_range)

if rescale_prenorm_residual:
# Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme:
# > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale
# > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers.
# > – GPT-2 :: https://openai.com/blog/better-language-models/
#
# Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py
for name, p in module.named_parameters():
if name in [“out_proj.weight”, “fc2.weight”]:
# Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block
# Following Pytorch init, except scale by 1/sqrt(2 * n_layer)
# We need to reinit p since this code could be called multiple times
# Having just p *= scale would repeatedly scale it down
nn.init.kaiming_uniform(p, a=math.sqrt(5))
with torch.no_grad():
p /= math.sqrt(n_residuals_per_layer * n_layer)


def segm_init_weights(m):
if isinstance(m, nn.Linear):
trunc_normal_(m.weight, std=0.02)
if isinstance(m, nn.Linear) and m.bias is not None:
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.Conv2d):
# NOTE conv was left to pytorch default in my original init
lecun_normal_(m.weight)
if m.bias is not None:
nn.init.zeros_(m.bias)
elif isinstance(m, (nn.LayerNorm, nn.GroupNorm, nn.BatchNorm2d)):
nn.init.zeros_(m.bias)
nn.init.ones_(m.weight)


class VisionMamba(nn.Module):
def init(self,
img_size=224,
patch_size=16,
stride=16,
depth=24, # 24个block
embed_dim=192,
channels=3,
num_classes=1000,
ssm_cfg=None,
drop_rate=0., # 神经元失活
drop_path_rate=0.1, # 层失活
norm_epsilon: float = 1e-5,
rms_norm: bool = False,
initializer_cfg=None,
fused_add_norm=False,
residual_in_fp32=False,
device=None,
dtype=None,
ft_seq_len=None,
pt_hw_seq_len=14,
if_bidirectional=False,
final_pool_type=‘none’,
if_abs_pos_embed=False,
if_rope=False, # 是否需要将某个维度打乱顺序,增强模型鲁棒性
if_rope_residual=False,
flip_img_sequences_ratio=-1.,
if_bimamba=False,
bimamba_type=“none”,
if_cls_token=False,
if_devide_out=False,
init_layer_scale=None,
use_double_cls_token=False,
use_middle_cls_token=False,
**kwargs):
factory_kwargs = {“device”: device, “dtype”: dtype}
# add factory_kwargs into kwargs
kwargs.update(factory_kwargs)
super().init()
self.residual_in_fp32 = residual_in_fp32
self.fused_add_norm = fused_add_norm
self.if_bidirectional = if_bidirectional
self.final_pool_type = final_pool_type
self.if_abs_pos_embed = if_abs_pos_embed
self.if_rope = if_rope
self.if_rope_residual = if_rope_residual
self.flip_img_sequences_ratio = flip_img_sequences_ratio
# self.if_cls_token = if_cls_token （此处与Mamba源代码不同）
self.if_cls_token = False
self.use_double_cls_token = use_double_cls_token # cls_tokens拼接到头和尾
self.use_middle_cls_token = use_middle_cls_token # cls_tokens拼接到中间
self.num_tokens = 1 if if_cls_token else 0

# pretrain parameters
self.num_classes = num_classes
self.d_model = self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models

self.patch_embed = PatchEmbed(
img_size=img_size, patch_size=patch_size, stride=stride, in_chans=channels, embed_dim=embed_dim)
num_patches = self.patch_embed.num_patches

if if_cls_token:
if use_double_cls_token: # 使用该模式，cls_tokens拼接在头和尾
self.cls_token_head = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
self.cls_token_tail = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
self.num_tokens = 2
else:
self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
# self.num_tokens = 1

if if_abs_pos_embed:
# pos_embed是加在（cls_tokens和num_patches）总体上的。
# self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, self.embed_dim))
# 此处与mamba不同,划分了x和z的位置编码
self.pos_embed_x = nn.Parameter(torch.zeros(1, 256, self.embed_dim))
self.pos_embed_z = nn.Parameter(torch.zeros(1, 64, self.embed_dim))
self.pos_drop = nn.Dropout(p=drop_rate) # 随机失活,增强鲁棒性
# 是否对position_embedding进行增强
if if_rope:
half_head_dim = embed_dim // 2
hw_seq_len = img_size // patch_size
self.rope = VisionRotaryEmbeddingFast(
dim=half_head_dim,
pt_seq_len=pt_hw_seq_len,
ft_seq_len=hw_seq_len
)

# 此处将head去掉了。
# self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
# settings = ws_settings.Settings()
# config_module = importlib.import_module(“lib.config.%s.config” % settings.script_name)
# cfg = config_module.cfg
# self.head = build_box_head(cfg, embed_dim)


# TODO: release this comment
# drop path rate ,随机失活
dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule
# import ipdb;ipdb.set_trace()
inter_dpr = [0.0] + dpr
self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0. else nn.Identity()
# transformer blocks
self.layers = nn.ModuleList(
[
create_block(
embed_dim,
ssm_cfg=ssm_cfg,
norm_epsilon=norm_epsilon,
rms_norm=rms_norm,
residual_in_fp32=residual_in_fp32,
fused_add_norm=fused_add_norm,
layer_idx=i,
if_bimamba=if_bimamba,
bimamba_type=bimamba_type,
drop_path=inter_dpr[i],
if_devide_out=if_devide_out,
init_layer_scale=init_layer_scale,
**factory_kwargs,
)
for i in range(depth)
]
)

# output head
self.norm_f = (nn.LayerNorm if not rms_norm else RMSNorm)(
embed_dim, eps=norm_epsilon, **factory_kwargs
)

# self.pre_logits = nn.Identity()

# original init
self.patch_embed.apply(segm_init_weights)
# 此处与mamba不同，问题在于没有执行head相关操作
# self.head.apply(segm_init_weights)
# 在pos_embed中也进行了模板和搜索的划分
if if_abs_pos_embed:
# trunc_normal_：保证权重范围落在正态分布之间，加快收敛速度
trunc_normal_(self.pos_embed_x, std=.02)
trunc_normal_(self.pos_embed_z, std=.02)

if if_cls_token:
if use_double_cls_token:
trunc_normal_(self.cls_token_head, std=.02)
trunc_normal_(self.cls_token_tail, std=.02)
else:
trunc_normal_(self.cls_token, std=.02)

# mamba init
self.apply(
partial(
_init_weights,
n_layer=depth,
**(initializer_cfg if initializer_cfg is not None else {}),
)
)


def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
return {
i: layer.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype, **kwargs)
for i, layer in enumerate(self.layers)
}

@torch.jit.ignore
def no_weight_decay(self):
return {“pos_embed”, “cls_token”, “dist_token”, “cls_token_head”, “cls_token_tail”}

@torch.jit.ignore()
def load_pretrained(self, checkpoint_path, prefix=“”):
_load_weights(self, checkpoint_path, prefix)
# 此处有所修改,将单个输入改为了两个输入（模板和搜索）
def forward_features(self, z, x, inference_params=None, if_random_cls_token_position=False, if_random_token_rank=False):
# taken from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
# with slight modifications to add the dist_token
# 1）将输入 X 切成小方块
# 将输入 Z 切成小方块
# print(‘x.shape:’,x.shape)
# print(‘z.shape:’,z.shape)
x = self.patch_embed(x) #x.shape = torch.Size([B, 3, 256, 256]) -> torch.Size([2, 256, 384])
z = self.patch_embed(z) #z.shape = torch.Size([B, 3, 128, 128]) -> torch.Size([2, 64, 384])
# 输出切块之后的大小, M表示切块之后的大小,B表示batch_size
B, M, _ = x.shape

if self.if_cls_token: # False
if self.use_double_cls_token:
cls_token_head = self.cls_token_head.expand(B, -1, -1)
cls_token_tail = self.cls_token_tail.expand(B, -1, -1)
token_position = [0, M + 1]
x = torch.cat((cls_token_head, x, cls_token_tail), dim=1)
M = x.shape[1]
else:
if self.use_middle_cls_token:
cls_token = self.cls_token.expand(B, -1, -1)
token_position = M // 2 # 中间位置
# add cls token in the middle
x = torch.cat((x[:, :token_position, :], cls_token, x[:, token_position:, :]), dim=1)
elif if_random_cls_token_position: # （随机位置）
cls_token = self.cls_token.expand(B, -1, -1)
token_position = random.randint(0, M)
x = torch.cat((x[:, :token_position, :], cls_token, x[:, token_position:, :]), dim=1)
print("token_position: ", token_position)
else:
cls_token = self.cls_token.expand(B, -1, -1) # stole cls_tokens impl from Phil Wang, thanks
token_position = 0 # 头插
x = torch.cat((cls_token, x), dim=1)
M = x.shape[1] # 对M进行更新


# 此处由于划分了模板和搜索,pos_embed也需要单独执行
# 上文中将patches和CLS_Tokens拼接上，然后现在将pos_embed拼接
if self.if_abs_pos_embed: # True
x = x + self.pos_embed_x # x = x + positon_embemding =torch.Size([B, 256, 384]) + torch.Size([1, 256, 384]) = torch.Size([B, 256, 384])
z = z + self.pos_embed_z # z = z + positon_embemding =torch.Size([B, 64, 384]) + torch.Size([1, 64, 384]) = torch.Size([B, 64, 384])
# 这里是将patches和pos_embed拼接上之后再进行模板特征和搜索特征的拼接
x = torch.cat((z, x), dim=1) # torch.Size([B, 320, 768])
x = self.pos_drop(x) # x.shape = torch.Size([B, 320, 384])
# 如果采用的随机序列生成索引
# 就需要采用shuffle_indices来重新生成索引
if if_random_token_rank: #False
# 生成随机 shuffle 索引
shuffle_indices = torch.randperm(M) # 生成随机数
# 随机的情况需要打印出来看看
if isinstance(token_position, list):
print("original value: ", x[0, token_position[0], 0], x[0, token_position[1], 0])
else:
print("original value: ", x[0, token_position, 0])
print("original token_position: ", token_position)

# 执行 shuffle 打乱顺序
x = x[:, shuffle_indices, :]
if isinstance(token_position, list):
# 找到 cls token 在 shuffle 之后的新位置
new_token_position = [torch.where(shuffle_indices == token_position[i])[0].item() for i in range(len(token_position))]
token_position = new_token_position
else:
# 找到 cls token 在 shuffle 之后的新位置
token_position = torch.where(shuffle_indices == token_position)[0].item()

if isinstance(token_position, list):
print("new value: ", x[0, token_position[0], 0], x[0, token_position[1], 0])
else:
print("new value: ", x[0, token_position, 0])
print("new token_position: ", token_position)
# print(‘x.shape:’,x.shape) # torch.Size([2, 320, 768])
# 反转操作
if_flip_img_sequences = False
if self.flip_img_sequences_ratio > 0 and (self.flip_img_sequences_ratio - random.random()) > 1e-5: # False
x = x.flip([1])
if_flip_img_sequences = True
# print(‘x.shape:’,x.shape) # x.shape: torch.Size([2, 320, 768])
# mamba impl
residual = None
hidden_states = x
# print(‘1:’,x.shape)
# 非双向操作
if not self.if_bidirectional: # True
for layer in self.layers:
if if_flip_img_sequences and self.if_rope: # False
hidden_states = hidden_states.flip([1])
if residual is not None:
residual = residual.flip([1])

# rope about
if self.if_rope: # False
hidden_states = self.rope(hidden_states)
if residual is not None and self.if_rope_residual:
residual = self.rope(residual)

if if_flip_img_sequences and self.if_rope: # False
hidden_states = hidden_states.flip([1])
if residual is not None:
residual = residual.flip([1])

hidden_states, residual = layer(
hidden_states, residual, inference_params=inference_params
)
# print(‘2:’,x.shape)
# 双向传播 （一个forward,一个backward）
else: # False
# get two layers in a single for-loop
for i in range(len(self.layers) // 2):
if self.if_rope:
hidden_states = self.rope(hidden_states)
if residual is not None and self.if_rope_residual:
residual = self.rope(residual)
# 前向传播结果 hidden_states_f；前向传播的残差 residual_f
# Forward_SSM （前向传播计算的是偶数 i * 2 层）
hidden_states_f, residual_f = self.layers[i * 2](
hidden_states, residual, inference_params=inference_params
)
# 后向传播（所以后向传播时，应该计算奇数 即 i * 2 + 1）

hidden_states_b, residual_b = self.layers[i * 2 + 1](
hidden_states.flip([1]), None if residual == None else residual.flip([1]), inference_params=inference_params
)
# 因为有一个前向和一个后向，所以需要将两者相加
# 训练时,后向传播采取了翻转,所以现在相加时应该将其翻转回来
hidden_states = hidden_states_f + hidden_states_b.flip([1])
residual = residual_f + residual_b.flip([1])
# print(‘3:’,x.shape)
# print(‘4:’,x.shape)
if not self.fused_add_norm: #False
if residual is None:
residual = hidden_states
else:
residual = residual + self.drop_path(hidden_states)
hidden_states = self.norm_f(residual.to(dtype=self.norm_f.weight.dtype))
else: #True
# Set prenorm=False here since we don’t need the residual
fused_add_norm_fn = rms_norm_fn if isinstance(self.norm_f, RMSNorm) else layer_norm_fn
hidden_states = fused_add_norm_fn( # hidden_states.shape = torch.Size([B, 320, 384])
self.drop_path(hidden_states),
self.norm_f.weight,
self.norm_f.bias,
eps=self.norm_f.eps,
residual=residual,
prenorm=False,
residual_in_fp32=self.residual_in_fp32,
)
# print(‘5:’,x.shape)
# return only cls token if it exists
if self.if_cls_token: #False
if self.use_double_cls_token:
return (hidden_states[:, token_position[0], :] + hidden_states[:, token_position[1], :]) / 2
else:
if self.use_middle_cls_token:
return hidden_states[:, token_position, :]
elif if_random_cls_token_position:
return hidden_states[:, token_position, :]
else:
return hidden_states[:, token_position, :]
print(‘5:’,x.shape)
# 最后池化操作
# 进入MLP时,VisionMamba的输出是hidden_states.shape = torch.Size([2, 320, 384]) → Batch_size, num_patches, embedding_dim
# MLP接受2个维度，所以需要采用池化的方法对Vmamba的输出进行处理
# hidden_states[:, -1, :] = torch.Size([2, 384])
if self.final_pool_type == ‘none’:
return hidden_states[:, -1, :]
elif self.final_pool_type == ‘mean’: #True
return hidden_states.mean(dim=1) #hidden_states.shape = torch.Size([2, 384])
elif self.final_pool_type == ‘max’:
return hidden_states
elif self.final_pool_type == ‘all’:
return hidden_states
else:
raise NotImplementedError



# 整个前向传播过程
def forward(self, x, return_features=False, inference_params=None, if_random_cls_token_position=False, if_random_token_rank=False):
x = self.forward_features(x, inference_params, if_random_cls_token_position=if_random_cls_token_position, if_random_token_rank=if_random_token_rank)
if return_features:
return x #此处返回mamba_encoder的特征
x = self.head(x)
if self.final_pool_type == ‘max’:
x = x.max(dim=1)[0]
return x


@register_model
def vim_tiny_patch16_224_bimambav2_final_pool_mean_abs_pos_embed_with_midclstok_div2(pretrained=False, **kwargs):
model = VisionMamba(
patch_size=16, embed_dim=192, depth=24, rms_norm=True, residual_in_fp32=True, fused_add_norm=True, final_pool_type=‘mean’, if_abs_pos_embed=True, if_rope=False, if_rope_residual=False, bimamba_type=“v2”, if_cls_token=True, if_devide_out=True, use_middle_cls_token=True, **kwargs)
model.default_cfg = _cfg()
if pretrained:
checkpoint = torch.hub.load_state_dict_from_url(
url=“to.do”,
map_location=“cpu”, check_hash=True
)
model.load_state_dict(checkpoint[“model”])
return model

@register_model
def vim_tiny_patch16_stride8_224_bimambav2_final_pool_mean_abs_pos_embed_with_midclstok_div2(pretrained=False, **kwargs):
model = VisionMamba(
patch_size=16, stride=8, embed_dim=192, depth=24, rms_norm=True, residual_in_fp32=True, fused_add_norm=True, final_pool_type=‘mean’, if_abs_pos_embed=True, if_rope=False, if_rope_residual=False, bimamba_type=“v2”, if_cls_token=True, if_devide_out=True, use_middle_cls_token=True, **kwargs)
model.default_cfg = _cfg()
if pretrained:
checkpoint = torch.hub.load_state_dict_from_url(
url=“to.do”,
map_location=“cpu”, check_hash=True
)
model.load_state_dict(checkpoint[“model”])
return model

@register_model
def vim_small_patch16_224_bimambav2_final_pool_mean_abs_pos_embed_with_midclstok_div2(pretrained=False, **kwargs):
model = VisionMamba(
patch_size=16, embed_dim=384, depth=24, rms_norm=True, residual_in_fp32=True, fused_add_norm=True, final_pool_type=‘all’, if_abs_pos_embed=True, if_rope=False, if_rope_residual=False, bimamba_type=“v2”, if_cls_token=True, if_devide_out=True, use_middle_cls_token=True, **kwargs)
model.default_cfg = _cfg()

if pretrained:
# checkpoint = torch.hub.load_state_dict_from_url(
# url=“to.do”,
# map_location=“cpu”, check_hash=True
# )
# model.load_state_dict(checkpoint[“model”])
checkpoint = torch.load(pretrained, map_location=“cpu”)
missing_keys, unexpected_keys = model.load_state_dict(checkpoint[“model”], strict=False)
print('Load pretrained model from: ’ + pretrained)
return model

@register_model
def vim_small_patch16_stride8_224_bimambav2_final_pool_mean_abs_pos_embed_with_midclstok_div2(pretrained=False, **kwargs):
model = VisionMamba(
patch_size=16, stride=8, embed_dim=384, depth=24, rms_norm=True, residual_in_fp32=True, fused_add_norm=True, final_pool_type=‘mean’, if_abs_pos_embed=True, if_rope=False, if_rope_residual=False, bimamba_type=“v2”, if_cls_token=True, if_devide_out=True, use_middle_cls_token=True, **kwargs)
model.default_cfg = _cfg()
if pretrained:
checkpoint = torch.hub.load_state_dict_from_url(
url=“to.do”,
map_location=“cpu”, check_hash=True
)
model.load_state_dict(checkpoint[“model”])
return model

####################### 本文加载预训练模型 ######################
################ 使用 torch.load 直接从本地路径加载预训练模型,使用strict=Flase进行load_state_dict，允许部分加载，可能会跳过或忽略某些不匹配的键。
################### 此处采用的池化方法是mean，一般统计全局信息场景
# @register_model
# def vim_small_patch16_224_bimambav2_final_pool_mean_abs_pos_embed_with_midclstok_div2(pretrained=False, **kwargs):
# model = VisionMamba(
# patch_size=16, embed_dim=384, depth=24, rms_norm=True, residual_in_fp32=True, fused_add_norm=True, final_pool_type=‘all’, if_abs_pos_embed=True, if_rope=False, if_rope_residual=False, bimamba_type=“v2”, if_cls_token=True, if_devide_out=True, use_middle_cls_token=True, **kwargs)
# model.default_cfg = _cfg()

# if pretrained:
# checkpoint = torch.load(pretrained, map_location=“cpu”)
# missing_keys, unexpected_keys = model.load_state_dict(checkpoint[“model”], strict=False)
# print('Load pretrained model from: ’ + pretrained)
# return model


####################### 使用torch.hub.load_state_action_from_url从url加载预训练模型,并使用check_hash=True确保下载文件的完整性，这里并未提供strict参数,通常意味着需要完全匹配模型的所有键。
# 池化为all，则是考虑整个序列的场景，处理复杂序列和任务需要的细粒度信息
# @register_model
# def vim_small_patch16_224_bimambav2_final_pool_mean_abs_pos_embed_with_midclstok_div2(pretrained=False, **kwargs):
# model = VisionMamba(
# patch_size=16, embed_dim=384, depth=24, rms_norm=True, residual_in_fp32=True, fused_add_norm=True, final_pool_type=‘mean’, if_abs_pos_embed=True, if_rope=False, if_rope_residual=False, bimamba_type=“v2”, if_cls_token=True, if_devide_out=True, use_middle_cls_token=True, **kwargs)
# model.default_cfg = _cfg()
# if pretrained:
# checkpoint = torch.hub.load_state_dict_from_url(
# url=“to.do”,
# map_location=“cpu”, check_hash=True
# )
# model.load_state_dict(checkpoint[“model”])
# return model中的forward_features是否有关于head的部分？还是仅仅做了特征提取？