mambaout.py 24 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736
  1. """
  2. MambaOut models for image classification.
  3. Some implementations are modified from:
  4. timm (https://github.com/rwightman/pytorch-image-models),
  5. MetaFormer (https://github.com/sail-sg/metaformer),
  6. InceptionNeXt (https://github.com/sail-sg/inceptionnext)
  7. """
  8. from collections import OrderedDict
  9. from typing import List, Optional, Tuple, Type, Union
  10. import torch
  11. from torch import nn
  12. from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
  13. from timm.layers import trunc_normal_, DropPath, calculate_drop_path_rates, LayerNorm, LayerScale, ClNormMlpClassifierHead, get_act_layer
  14. from ._builder import build_model_with_cfg
  15. from ._features import feature_take_indices
  16. from ._manipulate import checkpoint_seq
  17. from ._registry import register_model, generate_default_cfgs
  18. class Stem(nn.Module):
  19. r""" Code modified from InternImage:
  20. https://github.com/OpenGVLab/InternImage
  21. """
  22. def __init__(
  23. self,
  24. in_chs: int = 3,
  25. out_chs: int = 96,
  26. mid_norm: bool = True,
  27. act_layer: Type[nn.Module] = nn.GELU,
  28. norm_layer: Type[nn.Module] = LayerNorm,
  29. device=None,
  30. dtype=None,
  31. ):
  32. dd = {'device': device, 'dtype': dtype}
  33. super().__init__()
  34. self.conv1 = nn.Conv2d(
  35. in_chs,
  36. out_chs // 2,
  37. kernel_size=3,
  38. stride=2,
  39. padding=1,
  40. **dd,
  41. )
  42. self.norm1 = norm_layer(out_chs // 2, **dd) if mid_norm else None
  43. self.act = act_layer()
  44. self.conv2 = nn.Conv2d(
  45. out_chs // 2,
  46. out_chs,
  47. kernel_size=3,
  48. stride=2,
  49. padding=1,
  50. **dd,
  51. )
  52. self.norm2 = norm_layer(out_chs, **dd)
  53. def forward(self, x):
  54. x = self.conv1(x)
  55. if self.norm1 is not None:
  56. x = x.permute(0, 2, 3, 1)
  57. x = self.norm1(x)
  58. x = x.permute(0, 3, 1, 2)
  59. x = self.act(x)
  60. x = self.conv2(x)
  61. x = x.permute(0, 2, 3, 1)
  62. x = self.norm2(x)
  63. return x
  64. class DownsampleNormFirst(nn.Module):
  65. def __init__(
  66. self,
  67. in_chs: int = 96,
  68. out_chs: int = 198,
  69. norm_layer: Type[nn.Module] = LayerNorm,
  70. device=None,
  71. dtype=None,
  72. ):
  73. dd = {'device': device, 'dtype': dtype}
  74. super().__init__()
  75. self.norm = norm_layer(in_chs, **dd)
  76. self.conv = nn.Conv2d(
  77. in_chs,
  78. out_chs,
  79. kernel_size=3,
  80. stride=2,
  81. padding=1,
  82. **dd,
  83. )
  84. def forward(self, x):
  85. x = self.norm(x)
  86. x = x.permute(0, 3, 1, 2)
  87. x = self.conv(x)
  88. x = x.permute(0, 2, 3, 1)
  89. return x
  90. class Downsample(nn.Module):
  91. def __init__(
  92. self,
  93. in_chs: int = 96,
  94. out_chs: int = 198,
  95. norm_layer: Type[nn.Module] = LayerNorm,
  96. device=None,
  97. dtype=None,
  98. ):
  99. dd = {'device': device, 'dtype': dtype}
  100. super().__init__()
  101. self.conv = nn.Conv2d(
  102. in_chs,
  103. out_chs,
  104. kernel_size=3,
  105. stride=2,
  106. padding=1,
  107. **dd,
  108. )
  109. self.norm = norm_layer(out_chs, **dd)
  110. def forward(self, x):
  111. x = x.permute(0, 3, 1, 2)
  112. x = self.conv(x)
  113. x = x.permute(0, 2, 3, 1)
  114. x = self.norm(x)
  115. return x
  116. class MlpHead(nn.Module):
  117. """ MLP classification head
  118. """
  119. def __init__(
  120. self,
  121. in_features: int,
  122. num_classes: int = 1000,
  123. pool_type: str = 'avg',
  124. act_layer: Type[nn.Module] = nn.GELU,
  125. mlp_ratio: Optional[int] = 4,
  126. norm_layer: Type[nn.Module] = LayerNorm,
  127. drop_rate: float = 0.,
  128. bias: bool = True,
  129. device=None,
  130. dtype=None,
  131. ):
  132. dd = {'device': device, 'dtype': dtype}
  133. super().__init__()
  134. if mlp_ratio is not None:
  135. hidden_size = int(mlp_ratio * in_features)
  136. else:
  137. hidden_size = None
  138. self.pool_type = pool_type
  139. self.in_features = in_features
  140. self.hidden_size = hidden_size or in_features
  141. self.norm = norm_layer(in_features, **dd)
  142. if hidden_size:
  143. self.pre_logits = nn.Sequential(OrderedDict([
  144. ('fc', nn.Linear(in_features, hidden_size, **dd)),
  145. ('act', act_layer()),
  146. ('norm', norm_layer(hidden_size, **dd))
  147. ]))
  148. self.num_features = hidden_size
  149. else:
  150. self.num_features = in_features
  151. self.pre_logits = nn.Identity()
  152. self.fc = nn.Linear(self.num_features, num_classes, bias=bias, **dd) if num_classes > 0 else nn.Identity()
  153. self.head_dropout = nn.Dropout(drop_rate)
  154. def reset(self, num_classes: int, pool_type: Optional[str] = None, reset_other: bool = False):
  155. if pool_type is not None:
  156. self.pool_type = pool_type
  157. if reset_other:
  158. self.norm = nn.Identity()
  159. self.pre_logits = nn.Identity()
  160. self.num_features = self.in_features
  161. self.fc = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
  162. def forward(self, x, pre_logits: bool = False):
  163. if self.pool_type == 'avg':
  164. x = x.mean((1, 2))
  165. x = self.norm(x)
  166. x = self.pre_logits(x)
  167. x = self.head_dropout(x)
  168. if pre_logits:
  169. return x
  170. x = self.fc(x)
  171. return x
  172. class GatedConvBlock(nn.Module):
  173. r""" Our implementation of Gated CNN Block: https://arxiv.org/pdf/1612.08083
  174. Args:
  175. conv_ratio: control the number of channels to conduct depthwise convolution.
  176. Conduct convolution on partial channels can improve paraitcal efficiency.
  177. The idea of partial channels is from ShuffleNet V2 (https://arxiv.org/abs/1807.11164) and
  178. also used by InceptionNeXt (https://arxiv.org/abs/2303.16900) and FasterNet (https://arxiv.org/abs/2303.03667)
  179. """
  180. def __init__(
  181. self,
  182. dim: int,
  183. expansion_ratio: float = 8 / 3,
  184. kernel_size: int = 7,
  185. conv_ratio: float = 1.0,
  186. ls_init_value: Optional[float] = None,
  187. norm_layer: Type[nn.Module] = LayerNorm,
  188. act_layer: Type[nn.Module] = nn.GELU,
  189. drop_path: float = 0.,
  190. device=None,
  191. dtype=None,
  192. **kwargs
  193. ):
  194. dd = {'device': device, 'dtype': dtype}
  195. super().__init__()
  196. self.norm = norm_layer(dim, **dd)
  197. hidden = int(expansion_ratio * dim)
  198. self.fc1 = nn.Linear(dim, hidden * 2, **dd)
  199. self.act = act_layer()
  200. conv_channels = int(conv_ratio * dim)
  201. self.split_indices = (hidden, hidden - conv_channels, conv_channels)
  202. self.conv = nn.Conv2d(
  203. conv_channels,
  204. conv_channels,
  205. kernel_size=kernel_size,
  206. padding=kernel_size // 2,
  207. groups=conv_channels,
  208. **dd,
  209. )
  210. self.fc2 = nn.Linear(hidden, dim, **dd)
  211. self.ls = LayerScale(dim, **dd) if ls_init_value is not None else nn.Identity()
  212. self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
  213. def forward(self, x):
  214. shortcut = x # [B, H, W, C]
  215. x = self.norm(x)
  216. x = self.fc1(x)
  217. g, i, c = torch.split(x, self.split_indices, dim=-1)
  218. c = c.permute(0, 3, 1, 2) # [B, H, W, C] -> [B, C, H, W]
  219. c = self.conv(c)
  220. c = c.permute(0, 2, 3, 1) # [B, C, H, W] -> [B, H, W, C]
  221. x = self.fc2(self.act(g) * torch.cat((i, c), dim=-1))
  222. x = self.ls(x)
  223. x = self.drop_path(x)
  224. return x + shortcut
  225. class MambaOutStage(nn.Module):
  226. def __init__(
  227. self,
  228. dim: int,
  229. dim_out: Optional[int] = None,
  230. depth: int = 4,
  231. expansion_ratio: float = 8 / 3,
  232. kernel_size: int = 7,
  233. conv_ratio: float = 1.0,
  234. downsample: str = '',
  235. ls_init_value: Optional[float] = None,
  236. norm_layer: Type[nn.Module] = LayerNorm,
  237. act_layer: Type[nn.Module] = nn.GELU,
  238. drop_path: float = 0.,
  239. device=None,
  240. dtype=None,
  241. ):
  242. dd = {'device': device, 'dtype': dtype}
  243. super().__init__()
  244. dim_out = dim_out or dim
  245. self.grad_checkpointing = False
  246. if downsample == 'conv':
  247. self.downsample = Downsample(dim, dim_out, norm_layer=norm_layer, **dd)
  248. elif downsample == 'conv_nf':
  249. self.downsample = DownsampleNormFirst(dim, dim_out, norm_layer=norm_layer, **dd)
  250. else:
  251. assert dim == dim_out
  252. self.downsample = nn.Identity()
  253. self.blocks = nn.Sequential(*[
  254. GatedConvBlock(
  255. dim=dim_out,
  256. expansion_ratio=expansion_ratio,
  257. kernel_size=kernel_size,
  258. conv_ratio=conv_ratio,
  259. ls_init_value=ls_init_value,
  260. norm_layer=norm_layer,
  261. act_layer=act_layer,
  262. drop_path=drop_path[j] if isinstance(drop_path, (list, tuple)) else drop_path,
  263. **dd,
  264. )
  265. for j in range(depth)
  266. ])
  267. def forward(self, x):
  268. x = self.downsample(x)
  269. if self.grad_checkpointing and not torch.jit.is_scripting():
  270. x = checkpoint_seq(self.blocks, x)
  271. else:
  272. x = self.blocks(x)
  273. return x
  274. class MambaOut(nn.Module):
  275. r""" MetaFormer
  276. A PyTorch impl of : `MetaFormer Baselines for Vision` -
  277. https://arxiv.org/abs/2210.13452
  278. Args:
  279. in_chans (int): Number of input image channels. Default: 3.
  280. num_classes (int): Number of classes for classification head. Default: 1000.
  281. depths (list or tuple): Number of blocks at each stage. Default: [3, 3, 9, 3].
  282. dims (int): Feature dimension at each stage. Default: [96, 192, 384, 576].
  283. downsample_layers: (list or tuple): Downsampling layers before each stage.
  284. drop_path_rate (float): Stochastic depth rate. Default: 0.
  285. output_norm: norm before classifier head. Default: partial(nn.LayerNorm, eps=1e-6).
  286. head_fn: classification head. Default: nn.Linear.
  287. head_dropout (float): dropout for MLP classifier. Default: 0.
  288. """
  289. def __init__(
  290. self,
  291. in_chans: int = 3,
  292. num_classes: int = 1000,
  293. global_pool: str = 'avg',
  294. depths: Tuple[int, ...] = (3, 3, 9, 3),
  295. dims: Tuple[int, ...] = (96, 192, 384, 576),
  296. norm_layer: Type[nn.Module] = LayerNorm,
  297. act_layer: Type[nn.Module] = nn.GELU,
  298. conv_ratio: float = 1.0,
  299. expansion_ratio: float = 8/3,
  300. kernel_size: int = 7,
  301. stem_mid_norm: bool = True,
  302. ls_init_value: Optional[float] = None,
  303. downsample: str = 'conv',
  304. drop_path_rate: float = 0.,
  305. drop_rate: float = 0.,
  306. head_fn: str = 'default',
  307. device=None,
  308. dtype=None,
  309. ):
  310. super().__init__()
  311. dd = {'device': device, 'dtype': dtype}
  312. self.num_classes = num_classes
  313. self.drop_rate = drop_rate
  314. self.output_fmt = 'NHWC'
  315. if not isinstance(depths, (list, tuple)):
  316. depths = [depths] # it means the model has only one stage
  317. if not isinstance(dims, (list, tuple)):
  318. dims = [dims]
  319. act_layer = get_act_layer(act_layer)
  320. num_stage = len(depths)
  321. self.num_stage = num_stage
  322. self.feature_info = []
  323. self.stem = Stem(
  324. in_chans,
  325. dims[0],
  326. mid_norm=stem_mid_norm,
  327. act_layer=act_layer,
  328. norm_layer=norm_layer,
  329. **dd,
  330. )
  331. prev_dim = dims[0]
  332. dp_rates = calculate_drop_path_rates(drop_path_rate, depths, stagewise=True)
  333. cur = 0
  334. curr_stride = 4
  335. self.stages = nn.Sequential()
  336. for i in range(num_stage):
  337. dim = dims[i]
  338. stride = 2 if curr_stride == 2 or i > 0 else 1
  339. curr_stride *= stride
  340. stage = MambaOutStage(
  341. dim=prev_dim,
  342. dim_out=dim,
  343. depth=depths[i],
  344. kernel_size=kernel_size,
  345. conv_ratio=conv_ratio,
  346. expansion_ratio=expansion_ratio,
  347. downsample=downsample if i > 0 else '',
  348. ls_init_value=ls_init_value,
  349. norm_layer=norm_layer,
  350. act_layer=act_layer,
  351. drop_path=dp_rates[i],
  352. **dd,
  353. )
  354. self.stages.append(stage)
  355. prev_dim = dim
  356. # NOTE feature_info use currently assumes stage 0 == stride 1, rest are stride 2
  357. self.feature_info += [dict(num_chs=prev_dim, reduction=curr_stride, module=f'stages.{i}')]
  358. cur += depths[i]
  359. if head_fn == 'default':
  360. # specific to this model, unusual norm -> pool -> fc -> act -> norm -> fc combo
  361. self.head = MlpHead(
  362. prev_dim,
  363. num_classes,
  364. pool_type=global_pool,
  365. drop_rate=drop_rate,
  366. norm_layer=norm_layer,
  367. **dd,
  368. )
  369. else:
  370. # more typical norm -> pool -> fc -> act -> fc
  371. self.head = ClNormMlpClassifierHead(
  372. prev_dim,
  373. num_classes,
  374. hidden_size=int(prev_dim * 4),
  375. pool_type=global_pool,
  376. norm_layer=norm_layer,
  377. drop_rate=drop_rate,
  378. **dd,
  379. )
  380. self.num_features = prev_dim
  381. self.head_hidden_size = self.head.num_features
  382. self.apply(self._init_weights)
  383. def _init_weights(self, m):
  384. if isinstance(m, (nn.Conv2d, nn.Linear)):
  385. trunc_normal_(m.weight, std=.02)
  386. if m.bias is not None:
  387. nn.init.constant_(m.bias, 0)
  388. @torch.jit.ignore
  389. def group_matcher(self, coarse=False):
  390. return dict(
  391. stem=r'^stem',
  392. blocks=r'^stages\.(\d+)' if coarse else [
  393. (r'^stages\.(\d+)\.downsample', (0,)), # blocks
  394. (r'^stages\.(\d+)\.blocks\.(\d+)', None),
  395. ]
  396. )
  397. @torch.jit.ignore
  398. def set_grad_checkpointing(self, enable=True):
  399. for s in self.stages:
  400. s.grad_checkpointing = enable
  401. @torch.jit.ignore
  402. def get_classifier(self) -> nn.Module:
  403. return self.head.fc
  404. def reset_classifier(self, num_classes: int, global_pool: Optional[str] = None):
  405. self.num_classes = num_classes
  406. self.head.reset(num_classes, global_pool)
  407. def forward_intermediates(
  408. self,
  409. x: torch.Tensor,
  410. indices: Optional[Union[int, List[int]]] = None,
  411. norm: bool = False,
  412. stop_early: bool = False,
  413. output_fmt: str = 'NCHW',
  414. intermediates_only: bool = False,
  415. ) -> Union[List[torch.Tensor], Tuple[torch.Tensor, List[torch.Tensor]]]:
  416. """ Forward features that returns intermediates.
  417. Args:
  418. x: Input image tensor
  419. indices: Take last n blocks if int, all if None, select matching indices if sequence
  420. norm: Apply norm layer to compatible intermediates
  421. stop_early: Stop iterating over blocks when last desired intermediate hit
  422. output_fmt: Shape of intermediate feature outputs
  423. intermediates_only: Only return intermediate features
  424. Returns:
  425. """
  426. assert output_fmt in ('NCHW', 'NHWC'), 'Output format must be one of NCHW or NHWC.'
  427. channel_first = output_fmt == 'NCHW'
  428. intermediates = []
  429. take_indices, max_index = feature_take_indices(len(self.stages), indices)
  430. # forward pass
  431. x = self.stem(x)
  432. if torch.jit.is_scripting() or not stop_early: # can't slice blocks in torchscript
  433. stages = self.stages
  434. else:
  435. stages = self.stages[:max_index + 1]
  436. for feat_idx, stage in enumerate(stages):
  437. x = stage(x)
  438. if feat_idx in take_indices:
  439. intermediates.append(x)
  440. if channel_first:
  441. # reshape to BCHW output format
  442. intermediates = [y.permute(0, 3, 1, 2).contiguous() for y in intermediates]
  443. if intermediates_only:
  444. return intermediates
  445. return x, intermediates
  446. def prune_intermediate_layers(
  447. self,
  448. indices: Union[int, List[int]] = 1,
  449. prune_norm: bool = False,
  450. prune_head: bool = True,
  451. ):
  452. """ Prune layers not required for specified intermediates.
  453. """
  454. take_indices, max_index = feature_take_indices(len(self.stages), indices)
  455. self.stages = self.stages[:max_index + 1] # truncate blocks w/ stem as idx 0
  456. if prune_head:
  457. self.reset_classifier(0, '')
  458. return take_indices
  459. def forward_features(self, x):
  460. x = self.stem(x)
  461. x = self.stages(x)
  462. return x
  463. def forward_head(self, x, pre_logits: bool = False):
  464. x = self.head(x, pre_logits=pre_logits) if pre_logits else self.head(x)
  465. return x
  466. def forward(self, x):
  467. x = self.forward_features(x)
  468. x = self.forward_head(x)
  469. return x
  470. def checkpoint_filter_fn(state_dict, model):
  471. if 'model' in state_dict:
  472. state_dict = state_dict['model']
  473. if 'stem.conv1.weight' in state_dict:
  474. return state_dict
  475. import re
  476. out_dict = {}
  477. for k, v in state_dict.items():
  478. k = k.replace('downsample_layers.0.', 'stem.')
  479. k = re.sub(r'stages.([0-9]+).([0-9]+)', r'stages.\1.blocks.\2', k)
  480. k = re.sub(r'downsample_layers.([0-9]+)', r'stages.\1.downsample', k)
  481. # remap head names
  482. if k.startswith('norm.'):
  483. # this is moving to head since it's after the pooling
  484. k = k.replace('norm.', 'head.norm.')
  485. elif k.startswith('head.'):
  486. k = k.replace('head.fc1.', 'head.pre_logits.fc.')
  487. k = k.replace('head.norm.', 'head.pre_logits.norm.')
  488. k = k.replace('head.fc2.', 'head.fc.')
  489. out_dict[k] = v
  490. return out_dict
  491. def _cfg(url='', **kwargs):
  492. return {
  493. 'url': url,
  494. 'num_classes': 1000, 'input_size': (3, 224, 224), 'test_input_size': (3, 288, 288),
  495. 'pool_size': (7, 7), 'crop_pct': 1.0, 'interpolation': 'bicubic',
  496. 'mean': IMAGENET_DEFAULT_MEAN, 'std': IMAGENET_DEFAULT_STD,
  497. 'first_conv': 'stem.conv1', 'classifier': 'head.fc',
  498. 'license': 'apache-2.0',
  499. **kwargs
  500. }
  501. default_cfgs = generate_default_cfgs({
  502. # original weights
  503. 'mambaout_femto.in1k': _cfg(
  504. hf_hub_id='timm/'),
  505. 'mambaout_kobe.in1k': _cfg(
  506. hf_hub_id='timm/'),
  507. 'mambaout_tiny.in1k': _cfg(
  508. hf_hub_id='timm/'),
  509. 'mambaout_small.in1k': _cfg(
  510. hf_hub_id='timm/'),
  511. 'mambaout_base.in1k': _cfg(
  512. hf_hub_id='timm/'),
  513. # timm experiments below
  514. 'mambaout_small_rw.sw_e450_in1k': _cfg(
  515. hf_hub_id='timm/',
  516. ),
  517. 'mambaout_base_short_rw.sw_e500_in1k': _cfg(
  518. hf_hub_id='timm/',
  519. crop_pct=0.95, test_crop_pct=1.0,
  520. ),
  521. 'mambaout_base_tall_rw.sw_e500_in1k': _cfg(
  522. hf_hub_id='timm/',
  523. crop_pct=0.95, test_crop_pct=1.0,
  524. ),
  525. 'mambaout_base_wide_rw.sw_e500_in1k': _cfg(
  526. hf_hub_id='timm/',
  527. crop_pct=0.95, test_crop_pct=1.0,
  528. ),
  529. 'mambaout_base_plus_rw.sw_e150_in12k_ft_in1k': _cfg(
  530. hf_hub_id='timm/',
  531. ),
  532. 'mambaout_base_plus_rw.sw_e150_r384_in12k_ft_in1k': _cfg(
  533. hf_hub_id='timm/',
  534. input_size=(3, 384, 384), test_input_size=(3, 384, 384), crop_mode='squash', pool_size=(12, 12),
  535. ),
  536. 'mambaout_base_plus_rw.sw_e150_in12k': _cfg(
  537. hf_hub_id='timm/',
  538. num_classes=11821,
  539. ),
  540. 'test_mambaout': _cfg(input_size=(3, 160, 160), test_input_size=(3, 192, 192), pool_size=(5, 5)),
  541. })
  542. def _create_mambaout(variant, pretrained=False, **kwargs):
  543. model = build_model_with_cfg(
  544. MambaOut, variant, pretrained,
  545. pretrained_filter_fn=checkpoint_filter_fn,
  546. feature_cfg=dict(out_indices=(0, 1, 2, 3), flatten_sequential=True),
  547. **kwargs,
  548. )
  549. return model
  550. # a series of MambaOut models
  551. @register_model
  552. def mambaout_femto(pretrained=False, **kwargs):
  553. model_args = dict(depths=(3, 3, 9, 3), dims=(48, 96, 192, 288))
  554. return _create_mambaout('mambaout_femto', pretrained=pretrained, **dict(model_args, **kwargs))
  555. # Kobe Memorial Version with 24 Gated CNN blocks
  556. @register_model
  557. def mambaout_kobe(pretrained=False, **kwargs):
  558. model_args = dict(depths=[3, 3, 15, 3], dims=[48, 96, 192, 288])
  559. return _create_mambaout('mambaout_kobe', pretrained=pretrained, **dict(model_args, **kwargs))
  560. @register_model
  561. def mambaout_tiny(pretrained=False, **kwargs):
  562. model_args = dict(depths=[3, 3, 9, 3], dims=[96, 192, 384, 576])
  563. return _create_mambaout('mambaout_tiny', pretrained=pretrained, **dict(model_args, **kwargs))
  564. @register_model
  565. def mambaout_small(pretrained=False, **kwargs):
  566. model_args = dict(depths=[3, 4, 27, 3], dims=[96, 192, 384, 576])
  567. return _create_mambaout('mambaout_small', pretrained=pretrained, **dict(model_args, **kwargs))
  568. @register_model
  569. def mambaout_base(pretrained=False, **kwargs):
  570. model_args = dict(depths=[3, 4, 27, 3], dims=[128, 256, 512, 768])
  571. return _create_mambaout('mambaout_base', pretrained=pretrained, **dict(model_args, **kwargs))
  572. @register_model
  573. def mambaout_small_rw(pretrained=False, **kwargs):
  574. model_args = dict(
  575. depths=[3, 4, 27, 3],
  576. dims=[96, 192, 384, 576],
  577. stem_mid_norm=False,
  578. downsample='conv_nf',
  579. ls_init_value=1e-6,
  580. head_fn='norm_mlp',
  581. )
  582. return _create_mambaout('mambaout_small_rw', pretrained=pretrained, **dict(model_args, **kwargs))
  583. @register_model
  584. def mambaout_base_short_rw(pretrained=False, **kwargs):
  585. model_args = dict(
  586. depths=(3, 3, 25, 3),
  587. dims=(128, 256, 512, 768),
  588. expansion_ratio=3.0,
  589. conv_ratio=1.25,
  590. stem_mid_norm=False,
  591. downsample='conv_nf',
  592. ls_init_value=1e-6,
  593. head_fn='norm_mlp',
  594. )
  595. return _create_mambaout('mambaout_base_short_rw', pretrained=pretrained, **dict(model_args, **kwargs))
  596. @register_model
  597. def mambaout_base_tall_rw(pretrained=False, **kwargs):
  598. model_args = dict(
  599. depths=(3, 4, 30, 3),
  600. dims=(128, 256, 512, 768),
  601. expansion_ratio=2.5,
  602. conv_ratio=1.25,
  603. stem_mid_norm=False,
  604. downsample='conv_nf',
  605. ls_init_value=1e-6,
  606. head_fn='norm_mlp',
  607. )
  608. return _create_mambaout('mambaout_base_tall_rw', pretrained=pretrained, **dict(model_args, **kwargs))
  609. @register_model
  610. def mambaout_base_wide_rw(pretrained=False, **kwargs):
  611. model_args = dict(
  612. depths=(3, 4, 27, 3),
  613. dims=(128, 256, 512, 768),
  614. expansion_ratio=3.0,
  615. conv_ratio=1.5,
  616. stem_mid_norm=False,
  617. downsample='conv_nf',
  618. ls_init_value=1e-6,
  619. act_layer='silu',
  620. head_fn='norm_mlp',
  621. )
  622. return _create_mambaout('mambaout_base_wide_rw', pretrained=pretrained, **dict(model_args, **kwargs))
  623. @register_model
  624. def mambaout_base_plus_rw(pretrained=False, **kwargs):
  625. model_args = dict(
  626. depths=(3, 4, 30, 3),
  627. dims=(128, 256, 512, 768),
  628. expansion_ratio=3.0,
  629. conv_ratio=1.5,
  630. stem_mid_norm=False,
  631. downsample='conv_nf',
  632. ls_init_value=1e-6,
  633. act_layer='silu',
  634. head_fn='norm_mlp',
  635. )
  636. return _create_mambaout('mambaout_base_plus_rw', pretrained=pretrained, **dict(model_args, **kwargs))
  637. @register_model
  638. def test_mambaout(pretrained=False, **kwargs):
  639. model_args = dict(
  640. depths=(1, 1, 3, 1),
  641. dims=(16, 32, 48, 64),
  642. expansion_ratio=3,
  643. stem_mid_norm=False,
  644. downsample='conv_nf',
  645. ls_init_value=1e-4,
  646. act_layer='silu',
  647. head_fn='norm_mlp',
  648. )
  649. return _create_mambaout('test_mambaout', pretrained=pretrained, **dict(model_args, **kwargs))