""" FocalNet As described in `Focal Modulation Networks` - https://arxiv.org/abs/2203.11926 Significant modifications and refactoring from the original impl at https://github.com/microsoft/FocalNet This impl is/has: * fully convolutional, NCHW tensor layout throughout, seemed to have minimal performance impact but more flexible * re-ordered downsample / layer so that striding always at beginning of layer (stage) * no input size constraints or input resolution/H/W tracking through the model * torchscript fixed and a number of quirks cleaned up * feature extraction support via `features_only=True` """ # -------------------------------------------------------- # FocalNets -- Focal Modulation Networks # Copyright (c) 2022 Microsoft # Licensed under The MIT License [see LICENSE for details] # Written by Jianwei Yang (jianwyan@microsoft.com) # -------------------------------------------------------- from functools import partial import torch import torch.nn as nn import torch.utils.checkpoint as checkpoint from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD from timm.layers import Mlp, DropPath, LayerNorm2d, trunc_normal_, ClassifierHead, NormMlpClassifierHead from ._builder import build_model_with_cfg from ._manipulate import named_apply from ._registry import register_model __all__ = ['FocalNet'] class FocalModulation(nn.Module): def __init__( self, dim, focal_window, focal_level, focal_factor=2, bias=True, use_post_norm=False, normalize_modulator=False, proj_drop=0., norm_layer=LayerNorm2d, ): super().__init__() self.dim = dim self.focal_window = focal_window self.focal_level = focal_level self.focal_factor = focal_factor self.use_post_norm = use_post_norm self.normalize_modulator = normalize_modulator self.input_split = [dim, dim, self.focal_level + 1] self.f = nn.Conv2d(dim, 2 * dim + (self.focal_level + 1), kernel_size=1, bias=bias) self.h = nn.Conv2d(dim, dim, kernel_size=1, bias=bias) self.act = nn.GELU() self.proj = nn.Conv2d(dim, dim, kernel_size=1) self.proj_drop = nn.Dropout(proj_drop) self.focal_layers = nn.ModuleList() self.kernel_sizes = [] for k in range(self.focal_level): kernel_size = self.focal_factor * k + self.focal_window self.focal_layers.append(nn.Sequential( nn.Conv2d(dim, dim, kernel_size=kernel_size, groups=dim, padding=kernel_size // 2, bias=False), nn.GELU(), )) self.kernel_sizes.append(kernel_size) self.norm = norm_layer(dim) if self.use_post_norm else nn.Identity() def forward(self, x): """ Args: x: input features with shape of (B, H, W, C) """ C = x.shape[1] # pre linear projection x = self.f(x) q, ctx, gates = torch.split(x, self.input_split, 1) # context aggreation ctx_all = 0 for l, focal_layer in enumerate(self.focal_layers): ctx = focal_layer(ctx) ctx_all = ctx_all + ctx * gates[:, l:l + 1] ctx_global = self.act(ctx.mean((2, 3), keepdim=True)) ctx_all = ctx_all + ctx_global * gates[:, self.focal_level:] # normalize context if self.normalize_modulator: ctx_all = ctx_all / (self.focal_level + 1) # focal modulation x_out = q * self.h(ctx_all) x_out = self.norm(x_out) # post linear projection x_out = self.proj(x_out) x_out = self.proj_drop(x_out) return x_out class LayerScale2d(nn.Module): def __init__(self, dim, init_values=1e-5, inplace=False): super().__init__() self.inplace = inplace self.gamma = nn.Parameter(init_values * torch.ones(dim)) def forward(self, x): gamma = self.gamma.view(1, -1, 1, 1) return x.mul_(gamma) if self.inplace else x * gamma class FocalNetBlock(nn.Module): r""" Focal Modulation Network Block. Args: dim (int): Number of input channels. mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. proj_drop (float, optional): Dropout rate. Default: 0.0 drop_path (float, optional): Stochastic depth rate. Default: 0.0 act_layer (nn.Module, optional): Activation layer. Default: nn.GELU norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm focal_level (int): Number of focal levels. focal_window (int): Focal window size at first focal level layerscale_value (float): Initial layerscale value use_post_norm (bool): Whether to use layernorm after modulation """ def __init__( self, dim, mlp_ratio=4., focal_level=1, focal_window=3, use_post_norm=False, use_post_norm_in_modulation=False, normalize_modulator=False, layerscale_value=1e-4, proj_drop=0., drop_path=0., act_layer=nn.GELU, norm_layer=LayerNorm2d, ): super().__init__() self.dim = dim self.mlp_ratio = mlp_ratio self.focal_window = focal_window self.focal_level = focal_level self.use_post_norm = use_post_norm self.norm1 = norm_layer(dim) if not use_post_norm else nn.Identity() self.modulation = FocalModulation( dim, focal_window=focal_window, focal_level=self.focal_level, use_post_norm=use_post_norm_in_modulation, normalize_modulator=normalize_modulator, proj_drop=proj_drop, norm_layer=norm_layer, ) self.norm1_post = norm_layer(dim) if use_post_norm else nn.Identity() self.ls1 = LayerScale2d(dim, layerscale_value) if layerscale_value is not None else nn.Identity() self.drop_path1 = DropPath(drop_path) if drop_path > 0. else nn.Identity() self.norm2 = norm_layer(dim) if not use_post_norm else nn.Identity() self.mlp = Mlp( in_features=dim, hidden_features=int(dim * mlp_ratio), act_layer=act_layer, drop=proj_drop, use_conv=True, ) self.norm2_post = norm_layer(dim) if use_post_norm else nn.Identity() self.ls2 = LayerScale2d(dim, layerscale_value) if layerscale_value is not None else nn.Identity() self.drop_path2 = DropPath(drop_path) if drop_path > 0. else nn.Identity() def forward(self, x): shortcut = x # Focal Modulation x = self.norm1(x) x = self.modulation(x) x = self.norm1_post(x) x = shortcut + self.drop_path1(self.ls1(x)) # FFN x = x + self.drop_path2(self.ls2(self.norm2_post(self.mlp(self.norm2(x))))) return x class BasicLayer(nn.Module): """ A basic Focal Transformer layer for one stage. Args: dim (int): Number of input channels. depth (int): Number of blocks. mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. drop (float, optional): Dropout rate. Default: 0.0 drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0 norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm downsample (bool): Downsample layer at start of the layer. Default: True focal_level (int): Number of focal levels focal_window (int): Focal window size at first focal level layerscale_value (float): Initial layerscale value use_post_norm (bool): Whether to use layer norm after modulation """ def __init__( self, dim, out_dim, depth, mlp_ratio=4., downsample=True, focal_level=1, focal_window=1, use_overlap_down=False, use_post_norm=False, use_post_norm_in_modulation=False, normalize_modulator=False, layerscale_value=1e-4, proj_drop=0., drop_path=0., norm_layer=LayerNorm2d, ): super().__init__() self.dim = dim self.depth = depth self.grad_checkpointing = False if downsample: self.downsample = Downsample( in_chs=dim, out_chs=out_dim, stride=2, overlap=use_overlap_down, norm_layer=norm_layer, ) else: self.downsample = nn.Identity() # build blocks self.blocks = nn.ModuleList([ FocalNetBlock( dim=out_dim, mlp_ratio=mlp_ratio, focal_level=focal_level, focal_window=focal_window, use_post_norm=use_post_norm, use_post_norm_in_modulation=use_post_norm_in_modulation, normalize_modulator=normalize_modulator, layerscale_value=layerscale_value, proj_drop=proj_drop, drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path, norm_layer=norm_layer, ) for i in range(depth)]) @torch.jit.ignore def set_grad_checkpointing(self, enable=True): self.grad_checkpointing = enable def forward(self, x): x = self.downsample(x) for blk in self.blocks: if self.grad_checkpointing and not torch.jit.is_scripting(): x = checkpoint.checkpoint(blk, x) else: x = blk(x) return x class Downsample(nn.Module): r""" Args: in_chs (int): Number of input image channels out_chs (int): Number of linear projection output channels stride (int): Downsample stride. Default: 4. norm_layer (nn.Module, optional): Normalization layer. Default: None """ def __init__( self, in_chs, out_chs, stride=4, overlap=False, norm_layer=None, ): super().__init__() self.stride = stride padding = 0 kernel_size = stride if overlap: assert stride in (2, 4) if stride == 4: kernel_size, padding = 7, 2 elif stride == 2: kernel_size, padding = 3, 1 self.proj = nn.Conv2d(in_chs, out_chs, kernel_size=kernel_size, stride=stride, padding=padding) self.norm = norm_layer(out_chs) if norm_layer is not None else nn.Identity() def forward(self, x): x = self.proj(x) x = self.norm(x) return x class FocalNet(nn.Module): r""" Focal Modulation Networks (FocalNets) Args: in_chans (int): Number of input image channels. Default: 3 num_classes (int): Number of classes for classification head. Default: 1000 embed_dim (int): Patch embedding dimension. Default: 96 depths (tuple(int)): Depth of each Focal Transformer layer. mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4 focal_levels (list): How many focal levels at all stages. Note that this excludes the finest-grain level. Default: [1, 1, 1, 1] focal_windows (list): The focal window size at all stages. Default: [7, 5, 3, 1] use_overlap_down (bool): Whether to use convolutional embedding. use_post_norm (bool): Whether to use layernorm after modulation (it helps stablize training of large models) layerscale_value (float): Value for layer scale. Default: 1e-4 drop_rate (float): Dropout rate. Default: 0 drop_path_rate (float): Stochastic depth rate. Default: 0.1 norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm. """ def __init__( self, in_chans=3, num_classes=1000, global_pool='avg', embed_dim=96, depths=(2, 2, 6, 2), mlp_ratio=4., focal_levels=(2, 2, 2, 2), focal_windows=(3, 3, 3, 3), use_overlap_down=False, use_post_norm=False, use_post_norm_in_modulation=False, normalize_modulator=False, head_hidden_size=None, head_init_scale=1.0, layerscale_value=None, drop_rate=0., proj_drop_rate=0., drop_path_rate=0.1, norm_layer=partial(LayerNorm2d, eps=1e-5), **kwargs, ): super().__init__() self.num_layers = len(depths) embed_dim = [embed_dim * (2 ** i) for i in range(self.num_layers)] self.num_classes = num_classes self.embed_dim = embed_dim self.num_features = embed_dim[-1] self.feature_info = [] self.stem = Downsample( in_chs=in_chans, out_chs=embed_dim[0], overlap=use_overlap_down, norm_layer=norm_layer, ) in_dim = embed_dim[0] dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule layers = [] for i_layer in range(self.num_layers): out_dim = embed_dim[i_layer] layer = BasicLayer( dim=in_dim, out_dim=out_dim, depth=depths[i_layer], mlp_ratio=mlp_ratio, downsample=i_layer > 0, focal_level=focal_levels[i_layer], focal_window=focal_windows[i_layer], use_overlap_down=use_overlap_down, use_post_norm=use_post_norm, use_post_norm_in_modulation=use_post_norm_in_modulation, normalize_modulator=normalize_modulator, layerscale_value=layerscale_value, proj_drop=proj_drop_rate, drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])], norm_layer=norm_layer, ) in_dim = out_dim layers += [layer] self.feature_info += [dict(num_chs=out_dim, reduction=4 * 2 ** i_layer, module=f'layers.{i_layer}')] self.layers = nn.Sequential(*layers) if head_hidden_size: self.norm = nn.Identity() self.head = NormMlpClassifierHead( self.num_features, num_classes, hidden_size=head_hidden_size, pool_type=global_pool, drop_rate=drop_rate, norm_layer=norm_layer, ) else: self.norm = norm_layer(self.num_features) self.head = ClassifierHead( self.num_features, num_classes, pool_type=global_pool, drop_rate=drop_rate ) named_apply(partial(_init_weights, head_init_scale=head_init_scale), self) @torch.jit.ignore def no_weight_decay(self): return {''} @torch.jit.ignore def set_grad_checkpointing(self, enable=True): self.grad_checkpointing = enable for l in self.layers: l.set_grad_checkpointing(enable=enable) @torch.jit.ignore def get_classifier(self): return self.classifier.fc def reset_classifier(self, num_classes, global_pool=None): self.classifier.reset(num_classes, global_pool=global_pool) def forward_features(self, x): x = self.stem(x) x = self.layers(x) x = self.norm(x) return x def forward_head(self, x, pre_logits: bool = False): return self.head(x, pre_logits=pre_logits) def forward(self, x): x = self.forward_features(x) x = self.forward_head(x) return x def _init_weights(module, name=None, head_init_scale=1.0): if isinstance(module, nn.Conv2d): trunc_normal_(module.weight, std=.02) if module.bias is not None: nn.init.zeros_(module.bias) elif isinstance(module, nn.Linear): trunc_normal_(module.weight, std=.02) if module.bias is not None: nn.init.zeros_(module.bias) if name and 'head.fc' in name: module.weight.data.mul_(head_init_scale) module.bias.data.mul_(head_init_scale) def _cfg(url='', **kwargs): return { 'url': url, 'num_classes': 1000, 'input_size': (3, 224, 224), 'pool_size': None, 'crop_pct': .9, 'interpolation': 'bicubic', 'mean': IMAGENET_DEFAULT_MEAN, 'std': IMAGENET_DEFAULT_STD, 'first_conv': 'stem.proj', 'classifier': 'head.fc', **kwargs } default_cfgs = { "focalnet_tiny_srf": _cfg( url='https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_tiny_srf.pth'), "focalnet_small_srf": _cfg( url='https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_small_srf.pth'), "focalnet_base_srf": _cfg( url='https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_base_srf.pth'), "focalnet_tiny_lrf": _cfg( url='https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_tiny_lrf.pth'), "focalnet_small_lrf": _cfg( url='https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_small_lrf.pth'), "focalnet_base_lrf": _cfg( url='https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_base_lrf.pth'), "focalnet_large_fl3": _cfg( url='https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_large_lrf_384.pth', input_size=(3, 384, 384), crop_pct=1.0, num_classes=21842), "focalnet_large_fl4": _cfg( url='https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_large_lrf_384_fl4.pth', input_size=(3, 384, 384), crop_pct=1.0, num_classes=21842), "focalnet_xlarge_fl3": _cfg( url='https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_xlarge_lrf_384.pth', input_size=(3, 384, 384), crop_pct=1.0, num_classes=21842), "focalnet_xlarge_fl4": _cfg( url='https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_xlarge_lrf_384_fl4.pth', input_size=(3, 384, 384), crop_pct=1.0, num_classes=21842), "focalnet_huge_fl3": _cfg( url='https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_huge_lrf_224.pth', num_classes=0), "focalnet_huge_fl4": _cfg( url='https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_huge_lrf_224_fl4.pth', num_classes=0), } def checkpoint_filter_fn(state_dict, model: FocalNet): if 'stem.proj.weight' in state_dict: return import re out_dict = {} if 'model' in state_dict: state_dict = state_dict['model'] dest_dict = model.state_dict() for k, v in state_dict.items(): k = re.sub(r'gamma_([0-9])', r'ls\1.gamma', k) k = k.replace('patch_embed', 'stem') k = re.sub(r'layers.(\d+).downsample', lambda x: f'layers.{int(x.group(1)) + 1}.downsample', k) if 'norm' in k and k not in dest_dict: k = re.sub(r'norm([0-9])', r'norm\1_post', k) k = k.replace('ln.', 'norm.') k = k.replace('head', 'head.fc') if dest_dict[k].shape != v.shape: v = v.reshape(dest_dict[k].shape) out_dict[k] = v return out_dict def _create_focalnet(variant, pretrained=False, **kwargs): default_out_indices = tuple(i for i, _ in enumerate(kwargs.get('depths', (1, 1, 3, 1)))) out_indices = kwargs.pop('out_indices', default_out_indices) model = build_model_with_cfg( FocalNet, variant, pretrained, pretrained_filter_fn=checkpoint_filter_fn, feature_cfg=dict(flatten_sequential=True, out_indices=out_indices), **kwargs) return model @register_model def focalnet_tiny_srf(pretrained=False, **kwargs): model_kwargs = dict(depths=[2, 2, 6, 2], embed_dim=96, **kwargs) return _create_focalnet('focalnet_tiny_srf', pretrained=pretrained, **model_kwargs) @register_model def focalnet_small_srf(pretrained=False, **kwargs): model_kwargs = dict(depths=[2, 2, 18, 2], embed_dim=96, **kwargs) return _create_focalnet('focalnet_small_srf', pretrained=pretrained, **model_kwargs) @register_model def focalnet_base_srf(pretrained=False, **kwargs): model_kwargs = dict(depths=[2, 2, 18, 2], embed_dim=128, **kwargs) return _create_focalnet('focalnet_base_srf', pretrained=pretrained, **model_kwargs) @register_model def focalnet_tiny_lrf(pretrained=False, **kwargs): model_kwargs = dict(depths=[2, 2, 6, 2], embed_dim=96, focal_levels=[3, 3, 3, 3], **kwargs) return _create_focalnet('focalnet_tiny_lrf', pretrained=pretrained, **model_kwargs) @register_model def focalnet_small_lrf(pretrained=False, **kwargs): model_kwargs = dict(depths=[2, 2, 18, 2], embed_dim=96, focal_levels=[3, 3, 3, 3], **kwargs) return _create_focalnet('focalnet_small_lrf', pretrained=pretrained, **model_kwargs) @register_model def focalnet_base_lrf(pretrained=False, **kwargs): model_kwargs = dict(depths=[2, 2, 18, 2], embed_dim=128, focal_levels=[3, 3, 3, 3], **kwargs) return _create_focalnet('focalnet_base_lrf', pretrained=pretrained, **model_kwargs) # FocalNet large+ models @register_model def focalnet_large_fl3(pretrained=False, **kwargs): model_kwargs = dict( depths=[2, 2, 18, 2], embed_dim=192, focal_levels=[3, 3, 3, 3], focal_windows=[5] * 4, use_post_norm=True, use_overlap_down=True, layerscale_value=1e-4, **kwargs) return _create_focalnet('focalnet_large_fl3', pretrained=pretrained, **model_kwargs) @register_model def focalnet_large_fl4(pretrained=False, **kwargs): model_kwargs = dict( depths=[2, 2, 18, 2], embed_dim=192, focal_levels=[4, 4, 4, 4], use_post_norm=True, use_overlap_down=True, layerscale_value=1e-4, **kwargs) return _create_focalnet('focalnet_large_fl4', pretrained=pretrained, **model_kwargs) @register_model def focalnet_xlarge_fl3(pretrained=False, **kwargs): model_kwargs = dict( depths=[2, 2, 18, 2], embed_dim=256, focal_levels=[3, 3, 3, 3], focal_windows=[5] * 4, use_post_norm=True, use_overlap_down=True, layerscale_value=1e-4, **kwargs) return _create_focalnet('focalnet_xlarge_fl3', pretrained=pretrained, **model_kwargs) @register_model def focalnet_xlarge_fl4(pretrained=False, **kwargs): model_kwargs = dict( depths=[2, 2, 18, 2], embed_dim=256, focal_levels=[4, 4, 4, 4], use_post_norm=True, use_overlap_down=True, layerscale_value=1e-4, **kwargs) return _create_focalnet('focalnet_xlarge_fl4', pretrained=pretrained, **model_kwargs) @register_model def focalnet_huge_fl3(pretrained=False, **kwargs): model_kwargs = dict( depths=[2, 2, 18, 2], embed_dim=352, focal_levels=[3, 3, 3, 3], focal_windows=[3] * 4, use_post_norm=True, use_post_norm_in_modulation=True, use_overlap_down=True, layerscale_value=1e-4, **kwargs) return _create_focalnet('focalnet_huge_fl3', pretrained=pretrained, **model_kwargs) @register_model def focalnet_huge_fl4(pretrained=False, **kwargs): model_kwargs = dict( depths=[2, 2, 18, 2], embed_dim=352, focal_levels=[4, 4, 4, 4], use_post_norm=True, use_post_norm_in_modulation=True, use_overlap_down=True, layerscale_value=1e-4, **kwargs) return _create_focalnet('focalnet_huge_fl4', pretrained=pretrained, **model_kwargs)