focalnet.py 27 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740
  1. """ FocalNet
  2. As described in `Focal Modulation Networks` - https://arxiv.org/abs/2203.11926
  3. Significant modifications and refactoring from the original impl at https://github.com/microsoft/FocalNet
  4. This impl is/has:
  5. * fully convolutional, NCHW tensor layout throughout, seemed to have minimal performance impact but more flexible
  6. * re-ordered downsample / layer so that striding always at beginning of layer (stage)
  7. * no input size constraints or input resolution/H/W tracking through the model
  8. * torchscript fixed and a number of quirks cleaned up
  9. * feature extraction support via `features_only=True`
  10. """
  11. # --------------------------------------------------------
  12. # FocalNets -- Focal Modulation Networks
  13. # Copyright (c) 2022 Microsoft
  14. # Licensed under The MIT License [see LICENSE for details]
  15. # Written by Jianwei Yang (jianwyan@microsoft.com)
  16. # --------------------------------------------------------
  17. from functools import partial
  18. from typing import Callable, List, Optional, Tuple, Type, Union
  19. import torch
  20. import torch.nn as nn
  21. from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
  22. from timm.layers import (
  23. Mlp,
  24. DropPath,
  25. LayerNorm2d,
  26. LayerScale2d,
  27. trunc_normal_,
  28. ClassifierHead,
  29. NormMlpClassifierHead,
  30. calculate_drop_path_rates,
  31. )
  32. from ._builder import build_model_with_cfg
  33. from ._features import feature_take_indices
  34. from ._manipulate import named_apply, checkpoint
  35. from ._registry import generate_default_cfgs, register_model
  36. __all__ = ['FocalNet']
  37. class FocalModulation(nn.Module):
  38. def __init__(
  39. self,
  40. dim: int,
  41. focal_window: int,
  42. focal_level: int,
  43. focal_factor: int = 2,
  44. bias: bool = True,
  45. use_post_norm: bool = False,
  46. normalize_modulator: bool = False,
  47. proj_drop: float = 0.,
  48. norm_layer: Type[nn.Module] = LayerNorm2d,
  49. device=None,
  50. dtype=None,
  51. ):
  52. dd = {'device': device, 'dtype': dtype}
  53. super().__init__()
  54. self.dim = dim
  55. self.focal_window = focal_window
  56. self.focal_level = focal_level
  57. self.focal_factor = focal_factor
  58. self.use_post_norm = use_post_norm
  59. self.normalize_modulator = normalize_modulator
  60. self.input_split = [dim, dim, self.focal_level + 1]
  61. self.f = nn.Conv2d(dim, 2 * dim + (self.focal_level + 1), kernel_size=1, bias=bias, **dd)
  62. self.h = nn.Conv2d(dim, dim, kernel_size=1, bias=bias, **dd)
  63. self.act = nn.GELU()
  64. self.proj = nn.Conv2d(dim, dim, kernel_size=1, **dd)
  65. self.proj_drop = nn.Dropout(proj_drop)
  66. self.focal_layers = nn.ModuleList()
  67. self.kernel_sizes = []
  68. for k in range(self.focal_level):
  69. kernel_size = self.focal_factor * k + self.focal_window
  70. self.focal_layers.append(nn.Sequential(
  71. nn.Conv2d(dim, dim, kernel_size=kernel_size, groups=dim, padding=kernel_size // 2, bias=False, **dd),
  72. nn.GELU(),
  73. ))
  74. self.kernel_sizes.append(kernel_size)
  75. self.norm = norm_layer(dim, **dd) if self.use_post_norm else nn.Identity()
  76. def forward(self, x):
  77. # pre linear projection
  78. x = self.f(x)
  79. q, ctx, gates = torch.split(x, self.input_split, 1)
  80. # context aggregation
  81. ctx_all = 0
  82. for l, focal_layer in enumerate(self.focal_layers):
  83. ctx = focal_layer(ctx)
  84. ctx_all = ctx_all + ctx * gates[:, l:l + 1]
  85. ctx_global = self.act(ctx.mean((2, 3), keepdim=True))
  86. ctx_all = ctx_all + ctx_global * gates[:, self.focal_level:]
  87. # normalize context
  88. if self.normalize_modulator:
  89. ctx_all = ctx_all / (self.focal_level + 1)
  90. # focal modulation
  91. x_out = q * self.h(ctx_all)
  92. x_out = self.norm(x_out)
  93. # post linear projection
  94. x_out = self.proj(x_out)
  95. x_out = self.proj_drop(x_out)
  96. return x_out
  97. class FocalNetBlock(nn.Module):
  98. """ Focal Modulation Network Block.
  99. """
  100. def __init__(
  101. self,
  102. dim: int,
  103. mlp_ratio: float = 4.,
  104. focal_level: int = 1,
  105. focal_window: int = 3,
  106. use_post_norm: bool = False,
  107. use_post_norm_in_modulation: bool = False,
  108. normalize_modulator: bool = False,
  109. layerscale_value: Optional[float] = 1e-4,
  110. proj_drop: float = 0.,
  111. drop_path: float = 0.,
  112. act_layer: Type[nn.Module] = nn.GELU,
  113. norm_layer: Type[nn.Module] = LayerNorm2d,
  114. device=None,
  115. dtype=None,
  116. ):
  117. """
  118. Args:
  119. dim: Number of input channels.
  120. mlp_ratio: Ratio of mlp hidden dim to embedding dim.
  121. focal_level: Number of focal levels.
  122. focal_window: Focal window size at first focal level.
  123. use_post_norm: Whether to use layer norm after modulation.
  124. use_post_norm_in_modulation: Whether to use layer norm in modulation.
  125. layerscale_value: Initial layerscale value.
  126. proj_drop: Dropout rate.
  127. drop_path: Stochastic depth rate.
  128. act_layer: Activation layer.
  129. norm_layer: Normalization layer.
  130. """
  131. dd = {'device': device, 'dtype': dtype}
  132. super().__init__()
  133. self.dim = dim
  134. self.mlp_ratio = mlp_ratio
  135. self.focal_window = focal_window
  136. self.focal_level = focal_level
  137. self.use_post_norm = use_post_norm
  138. self.norm1 = norm_layer(dim, **dd) if not use_post_norm else nn.Identity()
  139. self.modulation = FocalModulation(
  140. dim,
  141. focal_window=focal_window,
  142. focal_level=self.focal_level,
  143. use_post_norm=use_post_norm_in_modulation,
  144. normalize_modulator=normalize_modulator,
  145. proj_drop=proj_drop,
  146. norm_layer=norm_layer,
  147. **dd,
  148. )
  149. self.norm1_post = norm_layer(dim, **dd) if use_post_norm else nn.Identity()
  150. self.ls1 = LayerScale2d(dim, layerscale_value, **dd) if layerscale_value is not None else nn.Identity()
  151. self.drop_path1 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
  152. self.norm2 = norm_layer(dim, **dd) if not use_post_norm else nn.Identity()
  153. self.mlp = Mlp(
  154. in_features=dim,
  155. hidden_features=int(dim * mlp_ratio),
  156. act_layer=act_layer,
  157. drop=proj_drop,
  158. use_conv=True,
  159. **dd,
  160. )
  161. self.norm2_post = norm_layer(dim, **dd) if use_post_norm else nn.Identity()
  162. self.ls2 = LayerScale2d(dim, layerscale_value, **dd) if layerscale_value is not None else nn.Identity()
  163. self.drop_path2 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
  164. def forward(self, x):
  165. shortcut = x
  166. # Focal Modulation
  167. x = self.norm1(x)
  168. x = self.modulation(x)
  169. x = self.norm1_post(x)
  170. x = shortcut + self.drop_path1(self.ls1(x))
  171. # FFN
  172. x = x + self.drop_path2(self.ls2(self.norm2_post(self.mlp(self.norm2(x)))))
  173. return x
  174. class FocalNetStage(nn.Module):
  175. """ A basic Focal Transformer layer for one stage.
  176. """
  177. def __init__(
  178. self,
  179. dim: int,
  180. out_dim: int,
  181. depth: int,
  182. mlp_ratio: float = 4.,
  183. downsample: bool = True,
  184. focal_level: int = 1,
  185. focal_window: int = 1,
  186. use_overlap_down: bool = False,
  187. use_post_norm: bool = False,
  188. use_post_norm_in_modulation: bool = False,
  189. normalize_modulator: bool = False,
  190. layerscale_value: Optional[float] = 1e-4,
  191. proj_drop: float = 0.,
  192. drop_path: Union[float, List[float]] = 0.,
  193. norm_layer: Type[nn.Module] = LayerNorm2d,
  194. device=None,
  195. dtype=None,
  196. ):
  197. """
  198. Args:
  199. dim: Number of input channels.
  200. out_dim: Number of output channels.
  201. depth: Number of blocks.
  202. mlp_ratio: Ratio of mlp hidden dim to embedding dim.
  203. downsample: Downsample layer at start of the layer.
  204. focal_level: Number of focal levels
  205. focal_window: Focal window size at first focal level
  206. use_overlap_down: User overlapped convolution in downsample layer.
  207. use_post_norm: Whether to use layer norm after modulation.
  208. use_post_norm_in_modulation: Whether to use layer norm in modulation.
  209. layerscale_value: Initial layerscale value
  210. proj_drop: Dropout rate for projections.
  211. drop_path: Stochastic depth rate.
  212. norm_layer: Normalization layer.
  213. """
  214. dd = {'device': device, 'dtype': dtype}
  215. super().__init__()
  216. self.dim = dim
  217. self.depth = depth
  218. self.grad_checkpointing = False
  219. if downsample:
  220. self.downsample = Downsample(
  221. in_chs=dim,
  222. out_chs=out_dim,
  223. stride=2,
  224. overlap=use_overlap_down,
  225. norm_layer=norm_layer,
  226. **dd,
  227. )
  228. else:
  229. self.downsample = nn.Identity()
  230. # build blocks
  231. self.blocks = nn.ModuleList([
  232. FocalNetBlock(
  233. dim=out_dim,
  234. mlp_ratio=mlp_ratio,
  235. focal_level=focal_level,
  236. focal_window=focal_window,
  237. use_post_norm=use_post_norm,
  238. use_post_norm_in_modulation=use_post_norm_in_modulation,
  239. normalize_modulator=normalize_modulator,
  240. layerscale_value=layerscale_value,
  241. proj_drop=proj_drop,
  242. drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
  243. norm_layer=norm_layer,
  244. **dd,
  245. )
  246. for i in range(depth)])
  247. @torch.jit.ignore
  248. def set_grad_checkpointing(self, enable=True):
  249. self.grad_checkpointing = enable
  250. def forward(self, x):
  251. x = self.downsample(x)
  252. for blk in self.blocks:
  253. if self.grad_checkpointing and not torch.jit.is_scripting():
  254. x = checkpoint(blk, x)
  255. else:
  256. x = blk(x)
  257. return x
  258. class Downsample(nn.Module):
  259. def __init__(
  260. self,
  261. in_chs: int,
  262. out_chs: int,
  263. stride: int = 4,
  264. overlap: bool = False,
  265. norm_layer: Optional[Type[nn.Module]] = None,
  266. device=None,
  267. dtype=None,
  268. ):
  269. """
  270. Args:
  271. in_chs: Number of input image channels.
  272. out_chs: Number of linear projection output channels.
  273. stride: Downsample stride.
  274. overlap: Use overlapping convolutions if True.
  275. norm_layer: Normalization layer.
  276. """
  277. dd = {'device': device, 'dtype': dtype}
  278. super().__init__()
  279. self.stride = stride
  280. padding = 0
  281. kernel_size = stride
  282. if overlap:
  283. assert stride in (2, 4)
  284. if stride == 4:
  285. kernel_size, padding = 7, 2
  286. elif stride == 2:
  287. kernel_size, padding = 3, 1
  288. self.proj = nn.Conv2d(in_chs, out_chs, kernel_size=kernel_size, stride=stride, padding=padding, **dd)
  289. self.norm = norm_layer(out_chs, **dd) if norm_layer is not None else nn.Identity()
  290. def forward(self, x):
  291. x = self.proj(x)
  292. x = self.norm(x)
  293. return x
  294. class FocalNet(nn.Module):
  295. """" Focal Modulation Networks (FocalNets)
  296. """
  297. def __init__(
  298. self,
  299. in_chans: int = 3,
  300. num_classes: int = 1000,
  301. global_pool: str = 'avg',
  302. embed_dim: int = 96,
  303. depths: Tuple[int, ...] = (2, 2, 6, 2),
  304. mlp_ratio: float = 4.,
  305. focal_levels: Tuple[int, ...] = (2, 2, 2, 2),
  306. focal_windows: Tuple[int, ...] = (3, 3, 3, 3),
  307. use_overlap_down: bool = False,
  308. use_post_norm: bool = False,
  309. use_post_norm_in_modulation: bool = False,
  310. normalize_modulator: bool = False,
  311. head_hidden_size: Optional[int] = None,
  312. head_init_scale: float = 1.0,
  313. layerscale_value: Optional[float] = None,
  314. drop_rate: float = 0.,
  315. proj_drop_rate: float = 0.,
  316. drop_path_rate: float = 0.1,
  317. norm_layer: Type[nn.Module] = partial(LayerNorm2d, eps=1e-5),
  318. device=None,
  319. dtype=None,
  320. ):
  321. """
  322. Args:
  323. in_chans: Number of input image channels.
  324. num_classes: Number of classes for classification head.
  325. embed_dim: Patch embedding dimension.
  326. depths: Depth of each Focal Transformer layer.
  327. mlp_ratio: Ratio of mlp hidden dim to embedding dim.
  328. focal_levels: How many focal levels at all stages. Note that this excludes the finest-grain level.
  329. focal_windows: The focal window size at all stages.
  330. use_overlap_down: Whether to use convolutional embedding.
  331. use_post_norm: Whether to use layernorm after modulation (it helps stabilize training of large models)
  332. layerscale_value: Value for layer scale.
  333. drop_rate: Dropout rate.
  334. drop_path_rate: Stochastic depth rate.
  335. norm_layer: Normalization layer.
  336. """
  337. super().__init__()
  338. dd = {'device': device, 'dtype': dtype}
  339. self.num_layers = len(depths)
  340. embed_dim = [embed_dim * (2 ** i) for i in range(self.num_layers)]
  341. self.num_classes = num_classes
  342. self.in_chans = in_chans
  343. self.embed_dim = embed_dim
  344. self.num_features = self.head_hidden_size = embed_dim[-1]
  345. self.feature_info = []
  346. self.stem = Downsample(
  347. in_chs=in_chans,
  348. out_chs=embed_dim[0],
  349. overlap=use_overlap_down,
  350. norm_layer=norm_layer,
  351. **dd,
  352. )
  353. in_dim = embed_dim[0]
  354. dpr = calculate_drop_path_rates(drop_path_rate, sum(depths)) # stochastic depth decay rule
  355. layers = []
  356. for i_layer in range(self.num_layers):
  357. out_dim = embed_dim[i_layer]
  358. layer = FocalNetStage(
  359. dim=in_dim,
  360. out_dim=out_dim,
  361. depth=depths[i_layer],
  362. mlp_ratio=mlp_ratio,
  363. downsample=i_layer > 0,
  364. focal_level=focal_levels[i_layer],
  365. focal_window=focal_windows[i_layer],
  366. use_overlap_down=use_overlap_down,
  367. use_post_norm=use_post_norm,
  368. use_post_norm_in_modulation=use_post_norm_in_modulation,
  369. normalize_modulator=normalize_modulator,
  370. layerscale_value=layerscale_value,
  371. proj_drop=proj_drop_rate,
  372. drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],
  373. norm_layer=norm_layer,
  374. **dd,
  375. )
  376. in_dim = out_dim
  377. layers += [layer]
  378. self.feature_info += [dict(num_chs=out_dim, reduction=4 * 2 ** i_layer, module=f'layers.{i_layer}')]
  379. self.layers = nn.Sequential(*layers)
  380. if head_hidden_size:
  381. self.norm = nn.Identity()
  382. self.head_hidden_size = head_hidden_size
  383. self.head = NormMlpClassifierHead(
  384. self.num_features,
  385. num_classes,
  386. hidden_size=head_hidden_size,
  387. pool_type=global_pool,
  388. drop_rate=drop_rate,
  389. norm_layer=norm_layer,
  390. **dd,
  391. )
  392. else:
  393. self.norm = norm_layer(self.num_features, **dd)
  394. self.head = ClassifierHead(
  395. self.num_features,
  396. num_classes,
  397. pool_type=global_pool,
  398. drop_rate=drop_rate,
  399. **dd,
  400. )
  401. named_apply(partial(_init_weights, head_init_scale=head_init_scale), self)
  402. @torch.jit.ignore
  403. def no_weight_decay(self):
  404. return {''}
  405. @torch.jit.ignore
  406. def group_matcher(self, coarse=False):
  407. return dict(
  408. stem=r'^stem',
  409. blocks=[
  410. (r'^layers\.(\d+)', None),
  411. (r'^norm', (99999,))
  412. ] if coarse else [
  413. (r'^layers\.(\d+).downsample', (0,)),
  414. (r'^layers\.(\d+)\.\w+\.(\d+)', None),
  415. (r'^norm', (99999,)),
  416. ]
  417. )
  418. @torch.jit.ignore
  419. def set_grad_checkpointing(self, enable=True):
  420. self.grad_checkpointing = enable
  421. for l in self.layers:
  422. l.set_grad_checkpointing(enable=enable)
  423. @torch.jit.ignore
  424. def get_classifier(self) -> nn.Module:
  425. return self.head.fc
  426. def reset_classifier(self, num_classes: int, global_pool: Optional[str] = None):
  427. self.num_classes = num_classes
  428. self.head.reset(num_classes, pool_type=global_pool)
  429. def forward_intermediates(
  430. self,
  431. x: torch.Tensor,
  432. indices: Optional[Union[int, List[int]]] = None,
  433. norm: bool = False,
  434. stop_early: bool = False,
  435. output_fmt: str = 'NCHW',
  436. intermediates_only: bool = False,
  437. ) -> Union[List[torch.Tensor], Tuple[torch.Tensor, List[torch.Tensor]]]:
  438. """ Forward features that returns intermediates.
  439. Args:
  440. x: Input image tensor
  441. indices: Take last n blocks if int, all if None, select matching indices if sequence
  442. norm: Apply norm layer to compatible intermediates
  443. stop_early: Stop iterating over blocks when last desired intermediate hit
  444. output_fmt: Shape of intermediate feature outputs
  445. intermediates_only: Only return intermediate features
  446. Returns:
  447. """
  448. assert output_fmt in ('NCHW',), 'Output shape must be NCHW.'
  449. intermediates = []
  450. take_indices, max_index = feature_take_indices(len(self.layers), indices)
  451. # forward pass
  452. x = self.stem(x)
  453. if torch.jit.is_scripting() or not stop_early: # can't slice blocks in torchscript
  454. stages = self.layers
  455. else:
  456. stages = self.layers[:max_index + 1]
  457. last_idx = len(self.layers) - 1
  458. for feat_idx, stage in enumerate(stages):
  459. x = stage(x)
  460. if feat_idx in take_indices:
  461. if norm and feat_idx == last_idx:
  462. x_inter = self.norm(x) # applying final norm to last intermediate
  463. else:
  464. x_inter = x
  465. intermediates.append(x_inter)
  466. if intermediates_only:
  467. return intermediates
  468. if feat_idx == last_idx:
  469. x = self.norm(x)
  470. return x, intermediates
  471. def prune_intermediate_layers(
  472. self,
  473. indices: Union[int, List[int]] = 1,
  474. prune_norm: bool = False,
  475. prune_head: bool = True,
  476. ):
  477. """ Prune layers not required for specified intermediates.
  478. """
  479. take_indices, max_index = feature_take_indices(len(self.layers), indices)
  480. self.layers = self.layers[:max_index + 1] # truncate blocks w/ stem as idx 0
  481. if prune_norm:
  482. self.norm = nn.Identity()
  483. if prune_head:
  484. self.reset_classifier(0, '')
  485. return take_indices
  486. def forward_features(self, x):
  487. x = self.stem(x)
  488. x = self.layers(x)
  489. x = self.norm(x)
  490. return x
  491. def forward_head(self, x, pre_logits: bool = False):
  492. return self.head(x, pre_logits=pre_logits) if pre_logits else self.head(x)
  493. def forward(self, x):
  494. x = self.forward_features(x)
  495. x = self.forward_head(x)
  496. return x
  497. def _init_weights(module, name=None, head_init_scale=1.0):
  498. if isinstance(module, nn.Conv2d):
  499. trunc_normal_(module.weight, std=.02)
  500. if module.bias is not None:
  501. nn.init.zeros_(module.bias)
  502. elif isinstance(module, nn.Linear):
  503. trunc_normal_(module.weight, std=.02)
  504. if module.bias is not None:
  505. nn.init.zeros_(module.bias)
  506. if name and 'head.fc' in name:
  507. module.weight.data.mul_(head_init_scale)
  508. module.bias.data.mul_(head_init_scale)
  509. def _cfg(url='', **kwargs):
  510. return {
  511. 'url': url,
  512. 'num_classes': 1000, 'input_size': (3, 224, 224), 'pool_size': (7, 7),
  513. 'crop_pct': .9, 'interpolation': 'bicubic',
  514. 'mean': IMAGENET_DEFAULT_MEAN, 'std': IMAGENET_DEFAULT_STD,
  515. 'first_conv': 'stem.proj', 'classifier': 'head.fc',
  516. 'license': 'mit', **kwargs
  517. }
  518. default_cfgs = generate_default_cfgs({
  519. "focalnet_tiny_srf.ms_in1k": _cfg(
  520. hf_hub_id='timm/'),
  521. "focalnet_small_srf.ms_in1k": _cfg(
  522. hf_hub_id='timm/'),
  523. "focalnet_base_srf.ms_in1k": _cfg(
  524. hf_hub_id='timm/'),
  525. "focalnet_tiny_lrf.ms_in1k": _cfg(
  526. hf_hub_id='timm/'),
  527. "focalnet_small_lrf.ms_in1k": _cfg(
  528. hf_hub_id='timm/'),
  529. "focalnet_base_lrf.ms_in1k": _cfg(
  530. hf_hub_id='timm/'),
  531. "focalnet_large_fl3.ms_in22k": _cfg(
  532. hf_hub_id='timm/',
  533. input_size=(3, 384, 384), pool_size=(12, 12), crop_pct=1.0, num_classes=21842),
  534. "focalnet_large_fl4.ms_in22k": _cfg(
  535. hf_hub_id='timm/',
  536. input_size=(3, 384, 384), pool_size=(12, 12), crop_pct=1.0, num_classes=21842),
  537. "focalnet_xlarge_fl3.ms_in22k": _cfg(
  538. hf_hub_id='timm/',
  539. input_size=(3, 384, 384), pool_size=(12, 12), crop_pct=1.0, num_classes=21842),
  540. "focalnet_xlarge_fl4.ms_in22k": _cfg(
  541. hf_hub_id='timm/',
  542. input_size=(3, 384, 384), pool_size=(12, 12), crop_pct=1.0, num_classes=21842),
  543. "focalnet_huge_fl3.ms_in22k": _cfg(
  544. hf_hub_id='timm/',
  545. num_classes=21842),
  546. "focalnet_huge_fl4.ms_in22k": _cfg(
  547. hf_hub_id='timm/',
  548. num_classes=0),
  549. })
  550. def checkpoint_filter_fn(state_dict, model: FocalNet):
  551. state_dict = state_dict.get('model', state_dict)
  552. if 'stem.proj.weight' in state_dict:
  553. return state_dict
  554. import re
  555. out_dict = {}
  556. dest_dict = model.state_dict()
  557. for k, v in state_dict.items():
  558. k = re.sub(r'gamma_([0-9])', r'ls\1.gamma', k)
  559. k = k.replace('patch_embed', 'stem')
  560. k = re.sub(r'layers.(\d+).downsample', lambda x: f'layers.{int(x.group(1)) + 1}.downsample', k)
  561. if 'norm' in k and k not in dest_dict:
  562. k = re.sub(r'norm([0-9])', r'norm\1_post', k)
  563. k = k.replace('ln.', 'norm.')
  564. k = k.replace('head', 'head.fc')
  565. if k in dest_dict and dest_dict[k].numel() == v.numel() and dest_dict[k].shape != v.shape:
  566. v = v.reshape(dest_dict[k].shape)
  567. out_dict[k] = v
  568. return out_dict
  569. def _create_focalnet(variant, pretrained=False, **kwargs):
  570. default_out_indices = tuple(i for i, _ in enumerate(kwargs.get('depths', (1, 1, 3, 1))))
  571. out_indices = kwargs.pop('out_indices', default_out_indices)
  572. model = build_model_with_cfg(
  573. FocalNet, variant, pretrained,
  574. pretrained_filter_fn=checkpoint_filter_fn,
  575. feature_cfg=dict(flatten_sequential=True, out_indices=out_indices),
  576. **kwargs)
  577. return model
  578. @register_model
  579. def focalnet_tiny_srf(pretrained=False, **kwargs) -> FocalNet:
  580. model_kwargs = dict(depths=[2, 2, 6, 2], embed_dim=96, **kwargs)
  581. return _create_focalnet('focalnet_tiny_srf', pretrained=pretrained, **model_kwargs)
  582. @register_model
  583. def focalnet_small_srf(pretrained=False, **kwargs) -> FocalNet:
  584. model_kwargs = dict(depths=[2, 2, 18, 2], embed_dim=96, **kwargs)
  585. return _create_focalnet('focalnet_small_srf', pretrained=pretrained, **model_kwargs)
  586. @register_model
  587. def focalnet_base_srf(pretrained=False, **kwargs) -> FocalNet:
  588. model_kwargs = dict(depths=[2, 2, 18, 2], embed_dim=128, **kwargs)
  589. return _create_focalnet('focalnet_base_srf', pretrained=pretrained, **model_kwargs)
  590. @register_model
  591. def focalnet_tiny_lrf(pretrained=False, **kwargs) -> FocalNet:
  592. model_kwargs = dict(depths=[2, 2, 6, 2], embed_dim=96, focal_levels=[3, 3, 3, 3], **kwargs)
  593. return _create_focalnet('focalnet_tiny_lrf', pretrained=pretrained, **model_kwargs)
  594. @register_model
  595. def focalnet_small_lrf(pretrained=False, **kwargs) -> FocalNet:
  596. model_kwargs = dict(depths=[2, 2, 18, 2], embed_dim=96, focal_levels=[3, 3, 3, 3], **kwargs)
  597. return _create_focalnet('focalnet_small_lrf', pretrained=pretrained, **model_kwargs)
  598. @register_model
  599. def focalnet_base_lrf(pretrained=False, **kwargs) -> FocalNet:
  600. model_kwargs = dict(depths=[2, 2, 18, 2], embed_dim=128, focal_levels=[3, 3, 3, 3], **kwargs)
  601. return _create_focalnet('focalnet_base_lrf', pretrained=pretrained, **model_kwargs)
  602. # FocalNet large+ models
  603. @register_model
  604. def focalnet_large_fl3(pretrained=False, **kwargs) -> FocalNet:
  605. model_kwargs = dict(
  606. depths=[2, 2, 18, 2], embed_dim=192, focal_levels=[3, 3, 3, 3], focal_windows=[5] * 4,
  607. use_post_norm=True, use_overlap_down=True, layerscale_value=1e-4, **kwargs)
  608. return _create_focalnet('focalnet_large_fl3', pretrained=pretrained, **model_kwargs)
  609. @register_model
  610. def focalnet_large_fl4(pretrained=False, **kwargs) -> FocalNet:
  611. model_kwargs = dict(
  612. depths=[2, 2, 18, 2], embed_dim=192, focal_levels=[4, 4, 4, 4],
  613. use_post_norm=True, use_overlap_down=True, layerscale_value=1e-4, **kwargs)
  614. return _create_focalnet('focalnet_large_fl4', pretrained=pretrained, **model_kwargs)
  615. @register_model
  616. def focalnet_xlarge_fl3(pretrained=False, **kwargs) -> FocalNet:
  617. model_kwargs = dict(
  618. depths=[2, 2, 18, 2], embed_dim=256, focal_levels=[3, 3, 3, 3], focal_windows=[5] * 4,
  619. use_post_norm=True, use_overlap_down=True, layerscale_value=1e-4, **kwargs)
  620. return _create_focalnet('focalnet_xlarge_fl3', pretrained=pretrained, **model_kwargs)
  621. @register_model
  622. def focalnet_xlarge_fl4(pretrained=False, **kwargs) -> FocalNet:
  623. model_kwargs = dict(
  624. depths=[2, 2, 18, 2], embed_dim=256, focal_levels=[4, 4, 4, 4],
  625. use_post_norm=True, use_overlap_down=True, layerscale_value=1e-4, **kwargs)
  626. return _create_focalnet('focalnet_xlarge_fl4', pretrained=pretrained, **model_kwargs)
  627. @register_model
  628. def focalnet_huge_fl3(pretrained=False, **kwargs) -> FocalNet:
  629. model_kwargs = dict(
  630. depths=[2, 2, 18, 2], embed_dim=352, focal_levels=[3, 3, 3, 3], focal_windows=[3] * 4,
  631. use_post_norm=True, use_post_norm_in_modulation=True, use_overlap_down=True, layerscale_value=1e-4, **kwargs)
  632. return _create_focalnet('focalnet_huge_fl3', pretrained=pretrained, **model_kwargs)
  633. @register_model
  634. def focalnet_huge_fl4(pretrained=False, **kwargs) -> FocalNet:
  635. model_kwargs = dict(
  636. depths=[2, 2, 18, 2], embed_dim=352, focal_levels=[4, 4, 4, 4],
  637. use_post_norm=True, use_post_norm_in_modulation=True, use_overlap_down=True, layerscale_value=1e-4, **kwargs)
  638. return _create_focalnet('focalnet_huge_fl4', pretrained=pretrained, **model_kwargs)