diff --git a/timm/models/__init__.py b/timm/models/__init__.py index ca4d1167..a8a66935 100644 --- a/timm/models/__init__.py +++ b/timm/models/__init__.py @@ -61,6 +61,7 @@ from .selecsls import * from .senet import * from .sequencer import * from .sknet import * +from .starnet import * from .swiftformer import * from .swin_transformer import * from .swin_transformer_v2 import * diff --git a/timm/models/starnet.py b/timm/models/starnet.py new file mode 100644 index 00000000..93b1e537 --- /dev/null +++ b/timm/models/starnet.py @@ -0,0 +1,344 @@ +""" +Implementation of Prof-of-Concept Network: StarNet. + +We make StarNet as simple as possible [to show the key contribution of element-wise multiplication]: + - like NO layer-scale in network design, + - and NO EMA during training, + - which would improve the performance further. + +Created by: Xu Ma (Email: ma.xu1@northeastern.edu) +Modified Date: Mar/29/2024 +""" +from typing import Any, Dict, List, Optional, Set, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD +from timm.layers import DropPath, SelectAdaptivePool2d, Linear, LayerType, trunc_normal_ +from ._builder import build_model_with_cfg +from ._features import feature_take_indices +from ._manipulate import checkpoint_seq +from ._registry import register_model, generate_default_cfgs + +__all__ = ['StarNet'] + + +class ConvBN(nn.Sequential): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int = 1, + stride: int = 1, + padding: int = 0, + with_bn: bool = True, + **kwargs + ): + super().__init__() + self.add_module('conv', nn.Conv2d( + in_channels, out_channels, kernel_size, stride=stride, padding=padding, **kwargs)) + if with_bn: + self.add_module('bn', nn.BatchNorm2d(out_channels)) + nn.init.constant_(self.bn.weight, 1) + nn.init.constant_(self.bn.bias, 0) + + +class Block(nn.Module): + def __init__( + self, + dim: int, + mlp_ratio: int = 3, + drop_path: float = 0., + act_layer: LayerType = nn.ReLU6, + ): + super().__init__() + self.dwconv = ConvBN(dim, dim, 7, 1, 3, groups=dim, with_bn=True) + self.f1 = ConvBN(dim, mlp_ratio * dim, 1, with_bn=False) + self.f2 = ConvBN(dim, mlp_ratio * dim, 1, with_bn=False) + self.g = ConvBN(mlp_ratio * dim, dim, 1, with_bn=True) + self.dwconv2 = ConvBN(dim, dim, 7, 1, 3, groups=dim, with_bn=False) + self.act = act_layer() + self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + residual = x + x = self.dwconv(x) + x1, x2 = self.f1(x), self.f2(x) + x = self.act(x1) * x2 + x = self.dwconv2(self.g(x)) + x = residual + self.drop_path(x) + return x + + +class StarNet(nn.Module): + def __init__( + self, + base_dim: int = 32, + depths: List[int] = [3, 3, 12, 5], + mlp_ratio: int = 4, + drop_rate: float = 0., + drop_path_rate: float = 0., + act_layer: LayerType = nn.ReLU6, + num_classes: int = 1000, + in_chans: int = 3, + global_pool: str = 'avg', + output_stride: int = 32, + **kwargs, + ): + super().__init__() + assert output_stride == 32 + self.num_classes = num_classes + self.drop_rate = drop_rate + self.grad_checkpointing = False + self.feature_info = [] + stem_chs = 32 + + # stem layer + self.stem = nn.Sequential( + ConvBN(in_chans, stem_chs, kernel_size=3, stride=2, padding=1), + act_layer(), + ) + prev_chs = stem_chs + + # build stages + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth + stages = [] + cur = 0 + for i_layer in range(len(depths)): + embed_dim = base_dim * 2 ** i_layer + down_sampler = ConvBN(prev_chs, embed_dim, 3, stride=2, padding=1) + blocks = [Block(embed_dim, mlp_ratio, dpr[cur + i], act_layer) for i in range(depths[i_layer])] + cur += depths[i_layer] + prev_chs = embed_dim + stages.append(nn.Sequential(down_sampler, *blocks)) + self.feature_info.append(dict( + num_chs=prev_chs, reduction=2**(i_layer+2), module=f'stages.{i_layer}')) + self.stages = nn.Sequential(*stages) + # head + self.num_features = self.head_hidden_size = prev_chs + self.norm = nn.BatchNorm2d(self.num_features) + self.global_pool = SelectAdaptivePool2d(pool_type=global_pool) + self.flatten = nn.Flatten(1) if global_pool else nn.Identity() # don't flatten if pooling disabled + self.head = Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity() + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, (nn.Linear, nn.Conv2d)): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.BatchNorm2d): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + + @torch.jit.ignore + def no_weight_decay(self) -> Set: + return set() + + @torch.jit.ignore + def group_matcher(self, coarse: bool = False) -> Dict[str, Any]: + matcher = dict( + stem=r'^stem\.\d+', + blocks=[(r'^stages\.(\d+)', None), (r'^norm', (99999,))] + ) + return matcher + + @torch.jit.ignore + def set_grad_checkpointing(self, enable: bool = True): + self.grad_checkpointing = enable + + @torch.jit.ignore + def get_classifier(self) -> nn.Module: + return self.head + + def reset_classifier(self, num_classes: int, global_pool: Optional[str] = None): + self.num_classes = num_classes + if global_pool is not None: + # NOTE: cannot meaningfully change pooling of efficient head after creation + self.global_pool = SelectAdaptivePool2d(pool_type=global_pool) + self.flatten = nn.Flatten(1) if global_pool else nn.Identity() # don't flatten if pooling disabled + self.head = Linear(self.head_hidden_size, num_classes) if num_classes > 0 else nn.Identity() + + def forward_intermediates( + self, + x: torch.Tensor, + indices: Optional[Union[int, List[int]]] = None, + norm: bool = False, + stop_early: bool = False, + output_fmt: str = 'NCHW', + intermediates_only: bool = False, + ) -> Union[List[torch.Tensor], Tuple[torch.Tensor, List[torch.Tensor]]]: + """ Forward features that returns intermediates. + + Args: + x: Input image tensor + indices: Take last n blocks if int, all if None, select matching indices if sequence + norm: Apply norm layer to compatible intermediates + stop_early: Stop iterating over blocks when last desired intermediate hit + output_fmt: Shape of intermediate feature outputs + intermediates_only: Only return intermediate features + Returns: + + """ + assert output_fmt in ('NCHW',), 'Output shape must be NCHW.' + intermediates = [] + take_indices, max_index = feature_take_indices(len(self.stages), indices) + last_idx = len(self.stages) - 1 + + # forward pass + x = self.stem(x) + if torch.jit.is_scripting() or not stop_early: # can't slice blocks in torchscript + stages = self.stages + else: + stages = self.stages[:max_index + 1] + + for feat_idx, stage in enumerate(stages): + x = stage(x) + if feat_idx in take_indices: + if norm and feat_idx == last_idx: + x_inter = self.norm(x) # applying final norm last intermediate + else: + x_inter = x + intermediates.append(x_inter) + + if intermediates_only: + return intermediates + + x = self.norm(x) + + return x, intermediates + + def prune_intermediate_layers( + self, + indices: Union[int, List[int]] = 1, + prune_norm: bool = False, + prune_head: bool = True, + ): + """ Prune layers not required for specified intermediates. + """ + take_indices, max_index = feature_take_indices(len(self.stages), indices) + self.stages = self.stages[:max_index + 1] # truncate blocks w/ stem as idx 0 + if prune_norm: + self.norm = nn.Identity() + if prune_head: + self.reset_classifier(0, '') + return take_indices + + def forward_features(self, x: torch.Tensor) -> torch.Tensor: + x = self.stem(x) + if self.grad_checkpointing and not torch.jit.is_scripting(): + x = checkpoint_seq(self.stages, x, flatten=True) + else: + x = self.stages(x) + x = self.norm(x) + return x + + def forward_head(self, x: torch.Tensor, pre_logits: bool = False) -> torch.Tensor: + x = self.global_pool(x) + x = self.flatten(x) + if self.drop_rate > 0.: + x = F.dropout(x, p=self.drop_rate, training=self.training) + return x if pre_logits else self.head(x) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.forward_features(x) + x = self.forward_head(x) + return x + + +def checkpoint_filter_fn(state_dict: Dict[str, torch.Tensor], model: nn.Module) -> Dict[str, torch.Tensor]: + if 'state_dict' in state_dict: + state_dict = state_dict['state_dict'] + out_dict = state_dict + return out_dict + + +def _cfg(url: str = '', **kwargs: Any) -> Dict[str, Any]: + return { + 'url': url, 'num_classes': 1000, 'input_size': (3, 224, 224), 'pool_size': (7, 7), + 'crop_pct': 0.875, 'interpolation': 'bicubic', + 'mean': IMAGENET_DEFAULT_MEAN, 'std': IMAGENET_DEFAULT_STD, + 'first_conv': 'stem.0.conv', 'classifier': 'head', + 'paper_ids': 'arXiv:2403.19967', + 'paper_name': 'Rewrite the Stars', + 'origin_url': 'https://github.com/ma-xu/Rewrite-the-Stars', + **kwargs + } + + +default_cfgs = generate_default_cfgs({ + 'starnet_s1.in1k': _cfg( + # hf_hub_id='timm/', + url='https://github.com/ma-xu/Rewrite-the-Stars/releases/download/checkpoints_v1/starnet_s1.pth.tar', + ), + 'starnet_s2.in1k': _cfg( + # hf_hub_id='timm/', + url='https://github.com/ma-xu/Rewrite-the-Stars/releases/download/checkpoints_v1/starnet_s2.pth.tar', + ), + 'starnet_s3.in1k': _cfg( + # hf_hub_id='timm/', + url='https://github.com/ma-xu/Rewrite-the-Stars/releases/download/checkpoints_v1/starnet_s3.pth.tar', + ), + 'starnet_s4.in1k': _cfg( + # hf_hub_id='timm/', + url='https://github.com/ma-xu/Rewrite-the-Stars/releases/download/checkpoints_v1/starnet_s4.pth.tar', + ), + 'starnet_s050.untrained': _cfg(), + 'starnet_s100.untrained': _cfg(), + 'starnet_s150.untrained': _cfg(), +}) + + +def _create_starnet(variant: str, pretrained: bool = False, **kwargs: Any) -> StarNet: + model = build_model_with_cfg( + StarNet, variant, pretrained, + pretrained_filter_fn=checkpoint_filter_fn, + feature_cfg=dict(out_indices=(0, 1, 2, 3), flatten_sequential=True), + **kwargs, + ) + return model + + +@register_model +def starnet_s1(pretrained: bool = False, **kwargs: Any) -> StarNet: + model_args = dict(base_dim=24, depths=[2, 2, 8, 3]) + return _create_starnet('starnet_s1', pretrained=pretrained, **dict(model_args, **kwargs)) + + +@register_model +def starnet_s2(pretrained: bool = False, **kwargs: Any) -> StarNet: + model_args = dict(base_dim=32, depths=[1, 2, 6, 2]) + return _create_starnet('starnet_s2', pretrained=pretrained, **dict(model_args, **kwargs)) + + +@register_model +def starnet_s3(pretrained: bool = False, **kwargs: Any) -> StarNet: + model_args = dict(base_dim=32, depths=[2, 2, 8, 4]) + return _create_starnet('starnet_s3', pretrained=pretrained, **dict(model_args, **kwargs)) + + +@register_model +def starnet_s4(pretrained: bool = False, **kwargs: Any) -> StarNet: + model_args = dict(base_dim=32, depths=[3, 3, 12, 5]) + return _create_starnet('starnet_s4', pretrained=pretrained, **dict(model_args, **kwargs)) + + +# very small networks # +@register_model +def starnet_s050(pretrained: bool = False, **kwargs: Any) -> StarNet: + model_args = dict(base_dim=16, depths=[1, 1, 3, 1], mlp_ratio=3) + return _create_starnet('starnet_s050', pretrained=pretrained, **dict(model_args, **kwargs)) + + +@register_model +def starnet_s100(pretrained: bool = False, **kwargs: Any) -> StarNet: + model_args = dict(base_dim=20, depths=[1, 2, 4, 1], mlp_ratio=4) + return _create_starnet('starnet_s100', pretrained=pretrained, **dict(model_args, **kwargs)) + + +@register_model +def starnet_s150(pretrained: bool = False, **kwargs: Any) -> StarNet: + model_args = dict(base_dim=24, depths=[1, 2, 4, 2], mlp_ratio=3) + return _create_starnet('starnet_s150', pretrained=pretrained, **dict(model_args, **kwargs)) \ No newline at end of file